Initial commit
This commit is contained in:
commit
7b60283303
1
.gitignore
vendored
Normal file
1
.gitignore
vendored
Normal file
@ -0,0 +1 @@
|
||||
unslop-cv
|
||||
688
cmd/unslop-cv.go
Normal file
688
cmd/unslop-cv.go
Normal file
@ -0,0 +1,688 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"image"
|
||||
"log"
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/rwcarlsen/goexif/exif"
|
||||
"gocv.io/x/gocv"
|
||||
)
|
||||
|
||||
type CameraProfile struct {
|
||||
Make string
|
||||
Model string
|
||||
FocalLength float64
|
||||
CameraMatrix gocv.Mat
|
||||
DistCoeffs gocv.Mat
|
||||
}
|
||||
|
||||
type PhotoCorrector struct {
|
||||
cameraProfiles map[string]CameraProfile
|
||||
defaultProfile CameraProfile
|
||||
}
|
||||
|
||||
func NewPhotoCorrector() *PhotoCorrector {
|
||||
pc := &PhotoCorrector{
|
||||
cameraProfiles: make(map[string]CameraProfile),
|
||||
}
|
||||
pc.loadCameraProfiles()
|
||||
return pc
|
||||
}
|
||||
|
||||
func (pc *PhotoCorrector) loadCameraProfiles() {
|
||||
pc.defaultProfile = pc.createDefaultProfile()
|
||||
|
||||
pc.cameraProfiles["Canon_EOS R5"] = pc.createCanonEOSR5Profile()
|
||||
pc.cameraProfiles["Nikon_D850"] = pc.createNikonD850Profile()
|
||||
pc.cameraProfiles["Sony_A7R IV"] = pc.createSonyA7R4Profile()
|
||||
}
|
||||
|
||||
func (pc *PhotoCorrector) createDefaultProfile() CameraProfile {
|
||||
fx := 1800.0
|
||||
fy := 1800.0
|
||||
cx := 2000.0
|
||||
cy := 1500.0
|
||||
|
||||
cameraMatrix := gocv.NewMatWithSize(3, 3, gocv.MatTypeCV64F)
|
||||
cameraMatrix.SetDoubleAt(0, 0, fx)
|
||||
cameraMatrix.SetDoubleAt(0, 1, 0)
|
||||
cameraMatrix.SetDoubleAt(0, 2, cx)
|
||||
cameraMatrix.SetDoubleAt(1, 0, 0)
|
||||
cameraMatrix.SetDoubleAt(1, 1, fy)
|
||||
cameraMatrix.SetDoubleAt(1, 2, cy)
|
||||
cameraMatrix.SetDoubleAt(2, 0, 0)
|
||||
cameraMatrix.SetDoubleAt(2, 1, 0)
|
||||
cameraMatrix.SetDoubleAt(2, 2, 1)
|
||||
|
||||
distCoeffs := gocv.NewMatWithSize(1, 5, gocv.MatTypeCV64F)
|
||||
distCoeffs.SetDoubleAt(0, 0, -0.3)
|
||||
distCoeffs.SetDoubleAt(0, 1, 0.15)
|
||||
distCoeffs.SetDoubleAt(0, 2, 0.0)
|
||||
distCoeffs.SetDoubleAt(0, 3, 0.0)
|
||||
distCoeffs.SetDoubleAt(0, 4, -0.05)
|
||||
|
||||
return CameraProfile{
|
||||
Make: "Generic",
|
||||
Model: "WideAngle",
|
||||
FocalLength: 16.0,
|
||||
CameraMatrix: cameraMatrix,
|
||||
DistCoeffs: distCoeffs,
|
||||
}
|
||||
}
|
||||
|
||||
func (pc *PhotoCorrector) createCanonEOSR5Profile() CameraProfile {
|
||||
fx := 2100.0
|
||||
fy := 2100.0
|
||||
cx := 2250.0
|
||||
cy := 1832.0
|
||||
|
||||
cameraMatrix := gocv.NewMatWithSize(3, 3, gocv.MatTypeCV64F)
|
||||
cameraMatrix.SetDoubleAt(0, 0, fx)
|
||||
cameraMatrix.SetDoubleAt(0, 1, 0)
|
||||
cameraMatrix.SetDoubleAt(0, 2, cx)
|
||||
cameraMatrix.SetDoubleAt(1, 0, 0)
|
||||
cameraMatrix.SetDoubleAt(1, 1, fy)
|
||||
cameraMatrix.SetDoubleAt(1, 2, cy)
|
||||
cameraMatrix.SetDoubleAt(2, 0, 0)
|
||||
cameraMatrix.SetDoubleAt(2, 1, 0)
|
||||
cameraMatrix.SetDoubleAt(2, 2, 1)
|
||||
|
||||
distCoeffs := gocv.NewMatWithSize(1, 5, gocv.MatTypeCV64F)
|
||||
distCoeffs.SetDoubleAt(0, 0, -0.25)
|
||||
distCoeffs.SetDoubleAt(0, 1, 0.12)
|
||||
distCoeffs.SetDoubleAt(0, 2, 0.001)
|
||||
distCoeffs.SetDoubleAt(0, 3, 0.001)
|
||||
distCoeffs.SetDoubleAt(0, 4, -0.02)
|
||||
|
||||
return CameraProfile{
|
||||
Make: "Canon",
|
||||
Model: "EOS R5",
|
||||
FocalLength: 15.0,
|
||||
CameraMatrix: cameraMatrix,
|
||||
DistCoeffs: distCoeffs,
|
||||
}
|
||||
}
|
||||
|
||||
func (pc *PhotoCorrector) createNikonD850Profile() CameraProfile {
|
||||
fx := 1950.0
|
||||
fy := 1950.0
|
||||
cx := 2259.0
|
||||
cy := 1752.0
|
||||
|
||||
cameraMatrix := gocv.NewMatWithSize(3, 3, gocv.MatTypeCV64F)
|
||||
cameraMatrix.SetDoubleAt(0, 0, fx)
|
||||
cameraMatrix.SetDoubleAt(0, 1, 0)
|
||||
cameraMatrix.SetDoubleAt(0, 2, cx)
|
||||
cameraMatrix.SetDoubleAt(1, 0, 0)
|
||||
cameraMatrix.SetDoubleAt(1, 1, fy)
|
||||
cameraMatrix.SetDoubleAt(1, 2, cy)
|
||||
cameraMatrix.SetDoubleAt(2, 0, 0)
|
||||
cameraMatrix.SetDoubleAt(2, 1, 0)
|
||||
cameraMatrix.SetDoubleAt(2, 2, 1)
|
||||
|
||||
distCoeffs := gocv.NewMatWithSize(1, 5, gocv.MatTypeCV64F)
|
||||
distCoeffs.SetDoubleAt(0, 0, -0.28)
|
||||
distCoeffs.SetDoubleAt(0, 1, 0.14)
|
||||
distCoeffs.SetDoubleAt(0, 2, 0.002)
|
||||
distCoeffs.SetDoubleAt(0, 3, 0.001)
|
||||
distCoeffs.SetDoubleAt(0, 4, -0.03)
|
||||
|
||||
return CameraProfile{
|
||||
Make: "Nikon",
|
||||
Model: "D850",
|
||||
FocalLength: 14.0,
|
||||
CameraMatrix: cameraMatrix,
|
||||
DistCoeffs: distCoeffs,
|
||||
}
|
||||
}
|
||||
|
||||
func (pc *PhotoCorrector) createSonyA7R4Profile() CameraProfile {
|
||||
fx := 2000.0
|
||||
fy := 2000.0
|
||||
cx := 4784.0
|
||||
cy := 3168.0
|
||||
|
||||
cameraMatrix := gocv.NewMatWithSize(3, 3, gocv.MatTypeCV64F)
|
||||
cameraMatrix.SetDoubleAt(0, 0, fx)
|
||||
cameraMatrix.SetDoubleAt(0, 1, 0)
|
||||
cameraMatrix.SetDoubleAt(0, 2, cx)
|
||||
cameraMatrix.SetDoubleAt(1, 0, 0)
|
||||
cameraMatrix.SetDoubleAt(1, 1, fy)
|
||||
cameraMatrix.SetDoubleAt(1, 2, cy)
|
||||
cameraMatrix.SetDoubleAt(2, 0, 0)
|
||||
cameraMatrix.SetDoubleAt(2, 1, 0)
|
||||
cameraMatrix.SetDoubleAt(2, 2, 1)
|
||||
|
||||
distCoeffs := gocv.NewMatWithSize(1, 5, gocv.MatTypeCV64F)
|
||||
distCoeffs.SetDoubleAt(0, 0, -0.22)
|
||||
distCoeffs.SetDoubleAt(0, 1, 0.10)
|
||||
distCoeffs.SetDoubleAt(0, 2, 0.0015)
|
||||
distCoeffs.SetDoubleAt(0, 3, 0.0012)
|
||||
distCoeffs.SetDoubleAt(0, 4, -0.015)
|
||||
|
||||
return CameraProfile{
|
||||
Make: "Sony",
|
||||
Model: "ILCE-7RM4",
|
||||
FocalLength: 16.0,
|
||||
CameraMatrix: cameraMatrix,
|
||||
DistCoeffs: distCoeffs,
|
||||
}
|
||||
}
|
||||
|
||||
type EXIFData struct {
|
||||
Make string
|
||||
Model string
|
||||
FocalLength float64
|
||||
Aperture float64
|
||||
ISO int
|
||||
ImageWidth int
|
||||
ImageHeight int
|
||||
LensModel string
|
||||
}
|
||||
|
||||
func extractEXIFData(filename string) (*EXIFData, error) {
|
||||
file, err := os.Open(filename)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer file.Close()
|
||||
|
||||
x, err := exif.Decode(file)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
data := &EXIFData{}
|
||||
|
||||
if make, err := x.Get(exif.Make); err == nil {
|
||||
data.Make = strings.TrimSpace(make.String())
|
||||
}
|
||||
if model, err := x.Get(exif.Model); err == nil {
|
||||
data.Model = strings.TrimSpace(model.String())
|
||||
}
|
||||
|
||||
if focal, err := x.Get(exif.FocalLength); err == nil {
|
||||
if focalVal, err := focal.Rat(0); err == nil {
|
||||
f, _ := focalVal.Float64()
|
||||
data.FocalLength = f
|
||||
}
|
||||
}
|
||||
|
||||
if aperture, err := x.Get(exif.FNumber); err == nil {
|
||||
if apertureVal, err := aperture.Rat(0); err == nil {
|
||||
f, _ := apertureVal.Float64()
|
||||
data.Aperture = f
|
||||
}
|
||||
}
|
||||
|
||||
if iso, err := x.Get(exif.ISOSpeedRatings); err == nil {
|
||||
if isoVal, err := iso.Int(0); err == nil {
|
||||
data.ISO = isoVal
|
||||
}
|
||||
}
|
||||
|
||||
if width, err := x.Get(exif.ImageWidth); err == nil {
|
||||
if widthVal, err := width.Int(0); err == nil {
|
||||
data.ImageWidth = widthVal
|
||||
}
|
||||
}
|
||||
if height, err := x.Get(exif.ImageLength); err == nil {
|
||||
if heightVal, err := height.Int(0); err == nil {
|
||||
data.ImageHeight = heightVal
|
||||
}
|
||||
}
|
||||
|
||||
return data, nil
|
||||
}
|
||||
|
||||
func (pc *PhotoCorrector) getCameraMatrix(exifData *EXIFData, width, height int) gocv.Mat {
|
||||
sensorWidth := 36.0
|
||||
if exifData.FocalLength > 0 {
|
||||
fx := (exifData.FocalLength / sensorWidth) * float64(width)
|
||||
fy := fx
|
||||
cx := float64(width) / 2.0
|
||||
cy := float64(height) / 2.0
|
||||
|
||||
cameraMatrix := gocv.NewMatWithSize(3, 3, gocv.MatTypeCV64F)
|
||||
cameraMatrix.SetDoubleAt(0, 0, fx)
|
||||
cameraMatrix.SetDoubleAt(0, 2, cx)
|
||||
cameraMatrix.SetDoubleAt(1, 1, fy)
|
||||
cameraMatrix.SetDoubleAt(1, 2, cy)
|
||||
cameraMatrix.SetDoubleAt(2, 2, 1.0)
|
||||
|
||||
return cameraMatrix
|
||||
}
|
||||
|
||||
return pc.defaultProfile.CameraMatrix.Clone()
|
||||
}
|
||||
|
||||
func (pc *PhotoCorrector) getDistortionCoefficients(exifData *EXIFData) gocv.Mat {
|
||||
profileKey := fmt.Sprintf("%s_%s", exifData.Make, exifData.Model)
|
||||
|
||||
if profile, exists := pc.cameraProfiles[profileKey]; exists {
|
||||
return profile.DistCoeffs.Clone()
|
||||
}
|
||||
|
||||
if exifData.FocalLength > 0 && exifData.FocalLength < 20 {
|
||||
|
||||
k1 := -0.2 + (20.0-exifData.FocalLength)*0.02
|
||||
k2 := 0.1 + (20.0-exifData.FocalLength)*0.01
|
||||
|
||||
distCoeffs := gocv.NewMatWithSize(1, 5, gocv.MatTypeCV64F)
|
||||
distCoeffs.SetDoubleAt(0, 0, k1)
|
||||
distCoeffs.SetDoubleAt(0, 1, k2)
|
||||
distCoeffs.SetDoubleAt(0, 2, 0.0)
|
||||
distCoeffs.SetDoubleAt(0, 3, 0.0)
|
||||
distCoeffs.SetDoubleAt(0, 4, 0.0)
|
||||
|
||||
return distCoeffs
|
||||
}
|
||||
|
||||
return pc.defaultProfile.DistCoeffs.Clone()
|
||||
}
|
||||
|
||||
func (pc *PhotoCorrector) correctLensDistortion(img gocv.Mat, exifData *EXIFData) gocv.Mat {
|
||||
height, width := img.Rows(), img.Cols()
|
||||
|
||||
cameraMatrix := pc.getCameraMatrix(exifData, width, height)
|
||||
distCoeffs := pc.getDistortionCoefficients(exifData)
|
||||
|
||||
defer cameraMatrix.Close()
|
||||
defer distCoeffs.Close()
|
||||
|
||||
corrected := gocv.NewMat()
|
||||
|
||||
gocv.Undistort(img, &corrected, cameraMatrix, distCoeffs, cameraMatrix)
|
||||
|
||||
return corrected
|
||||
}
|
||||
|
||||
func (pc *PhotoCorrector) detectVanishingPoints(img gocv.Mat) []gocv.Point2f {
|
||||
gray := gocv.NewMat()
|
||||
defer gray.Close()
|
||||
gocv.CvtColor(img, &gray, gocv.ColorBGRToGray)
|
||||
|
||||
edges := gocv.NewMat()
|
||||
defer edges.Close()
|
||||
gocv.Canny(gray, &edges, 50, 150)
|
||||
|
||||
lines := gocv.NewMat()
|
||||
defer lines.Close()
|
||||
gocv.HoughLines(edges, &lines, 1, math.Pi/180, 100)
|
||||
|
||||
vanishingPoints := make([]gocv.Point2f, 0)
|
||||
|
||||
if lines.Rows() > 0 {
|
||||
|
||||
vp1 := gocv.Point2f{X: float32(img.Cols() / 2), Y: float32(img.Rows() / 2)}
|
||||
vanishingPoints = append(vanishingPoints, vp1)
|
||||
}
|
||||
|
||||
return vanishingPoints
|
||||
}
|
||||
|
||||
func (pc *PhotoCorrector) calculatePerspectiveCorrectionStrength(focalLength float64) float64 {
|
||||
if focalLength <= 0 {
|
||||
return 1.0
|
||||
}
|
||||
|
||||
switch {
|
||||
case focalLength < 16:
|
||||
return 1.5
|
||||
case focalLength < 24:
|
||||
return 1.2
|
||||
case focalLength < 35:
|
||||
return 1.0
|
||||
default:
|
||||
return 0.5
|
||||
}
|
||||
}
|
||||
|
||||
func (pc *PhotoCorrector) correctPerspective(img gocv.Mat, exifData *EXIFData) gocv.Mat {
|
||||
height, width := img.Rows(), img.Cols()
|
||||
|
||||
correctionStrength := pc.calculatePerspectiveCorrectionStrength(exifData.FocalLength)
|
||||
|
||||
if exifData.FocalLength > 85 {
|
||||
return img.Clone()
|
||||
}
|
||||
|
||||
vanishingPoints := pc.detectVanishingPoints(img)
|
||||
|
||||
if len(vanishingPoints) == 0 {
|
||||
return img.Clone()
|
||||
}
|
||||
|
||||
srcPoints := []image.Point{
|
||||
{X: 0, Y: 0},
|
||||
{X: width, Y: 0},
|
||||
{X: width, Y: height},
|
||||
{X: 0, Y: height},
|
||||
}
|
||||
|
||||
margin := int(math.Round(float64(width) * 0.05 * correctionStrength))
|
||||
dstPoints := []image.Point{
|
||||
{X: margin, Y: margin},
|
||||
{X: width - margin, Y: margin},
|
||||
{X: width - margin, Y: height - margin},
|
||||
{X: margin, Y: height - margin},
|
||||
}
|
||||
|
||||
transformMatrix := gocv.GetPerspectiveTransform(gocv.NewPointVectorFromPoints(srcPoints), gocv.NewPointVectorFromPoints(dstPoints))
|
||||
defer transformMatrix.Close()
|
||||
|
||||
size := img.Size()
|
||||
imageSize := image.Pt(size[1], size[0])
|
||||
|
||||
corrected := gocv.NewMat()
|
||||
gocv.WarpPerspective(img, &corrected, transformMatrix, imageSize)
|
||||
|
||||
return corrected
|
||||
}
|
||||
|
||||
func (pc *PhotoCorrector) enhanceForRealEstate(img gocv.Mat, exifData *EXIFData) gocv.Mat {
|
||||
enhanced := img.Clone()
|
||||
|
||||
if exifData.Aperture > 0 && exifData.ISO > 0 {
|
||||
|
||||
alpha := float32(1.0)
|
||||
beta := float32(0.0)
|
||||
|
||||
if exifData.ISO > 800 {
|
||||
alpha = 1.1
|
||||
}
|
||||
if exifData.Aperture < 2.8 {
|
||||
beta = 10.0
|
||||
}
|
||||
|
||||
tempMat := gocv.NewMat()
|
||||
enhanced.ConvertTo(&tempMat, gocv.MatTypeCV8U)
|
||||
tempMat.MultiplyFloat(alpha)
|
||||
tempMat.AddFloat(beta)
|
||||
enhanced = tempMat
|
||||
}
|
||||
|
||||
hsv := gocv.NewMat()
|
||||
defer hsv.Close()
|
||||
gocv.CvtColor(enhanced, &hsv, gocv.ColorBGRToHSV)
|
||||
|
||||
channels := gocv.Split(hsv)
|
||||
defer func() {
|
||||
for _, ch := range channels {
|
||||
ch.Close()
|
||||
}
|
||||
}()
|
||||
|
||||
tempSat := gocv.NewMat()
|
||||
channels[1].ConvertTo(&tempSat, gocv.MatTypeCV8U)
|
||||
tempSat.MultiplyFloat(1.1)
|
||||
channels[1] = tempSat
|
||||
|
||||
gocv.Merge(channels, &hsv)
|
||||
gocv.CvtColor(hsv, &enhanced, gocv.ColorHSVToBGR)
|
||||
|
||||
return enhanced
|
||||
}
|
||||
|
||||
func (pc *PhotoCorrector) ProcessImage(inputPath, outputPath string) error {
|
||||
img := gocv.IMRead(inputPath, gocv.IMReadColor)
|
||||
if img.Empty() {
|
||||
return fmt.Errorf("failed to load image: %s", inputPath)
|
||||
}
|
||||
defer img.Close()
|
||||
|
||||
exifData, err := extractEXIFData(inputPath)
|
||||
if err != nil {
|
||||
log.Printf("Warning: Could not extract EXIF data: %v", err)
|
||||
|
||||
exifData = &EXIFData{
|
||||
Make: "Unknown",
|
||||
Model: "Unknown",
|
||||
FocalLength: 16.0,
|
||||
ImageWidth: img.Cols(),
|
||||
ImageHeight: img.Rows(),
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Printf("Processing image with EXIF data: %+v\n", exifData)
|
||||
|
||||
fmt.Println("Correcting lens distortion...")
|
||||
distortionCorrected := pc.correctLensDistortion(img, exifData)
|
||||
defer distortionCorrected.Close()
|
||||
|
||||
fmt.Println("Correcting perspective...")
|
||||
perspectiveCorrected := pc.correctPerspective(distortionCorrected, exifData)
|
||||
defer perspectiveCorrected.Close()
|
||||
|
||||
fmt.Println("Applying real estate enhancements...")
|
||||
enhanced := pc.enhanceForRealEstate(perspectiveCorrected, exifData)
|
||||
defer enhanced.Close()
|
||||
|
||||
ok := gocv.IMWrite(outputPath, enhanced)
|
||||
if !ok {
|
||||
return fmt.Errorf("failed to save corrected image: %s", outputPath)
|
||||
}
|
||||
|
||||
fmt.Printf("Successfully processed image: %s -> %s\n", inputPath, outputPath)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (pc *PhotoCorrector) CalibrateCamera(calibrationImages []string, boardSize image.Point) (*CameraProfile, error) {
|
||||
if len(calibrationImages) == 0 {
|
||||
return nil, fmt.Errorf("no calibration images provided")
|
||||
}
|
||||
|
||||
objectPoints := make([][]gocv.Point3f, 0)
|
||||
objectPointsVector := gocv.NewPoints3fVectorFromPoints(objectPoints)
|
||||
imagePoints := make([][]gocv.Point2f, 0)
|
||||
imagePointsVector := gocv.NewPoints2fVectorFromPoints(imagePoints)
|
||||
|
||||
objp := make([]gocv.Point3f, 0)
|
||||
for i := range boardSize.Y {
|
||||
for j := range boardSize.X {
|
||||
objp = append(objp, gocv.Point3f{
|
||||
X: float32(j), Y: float32(i), Z: 0,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
var imageSize image.Point
|
||||
|
||||
fmt.Printf("Processing %d calibration images...\n", len(calibrationImages))
|
||||
for i, imagePath := range calibrationImages {
|
||||
img := gocv.IMRead(imagePath, gocv.IMReadColor)
|
||||
if img.Empty() {
|
||||
fmt.Printf("Warning: Could not load image %s\n", imagePath)
|
||||
continue
|
||||
}
|
||||
|
||||
if imageSize.X == 0 {
|
||||
size := img.Size()
|
||||
imageSize = image.Pt(size[1], size[0])
|
||||
}
|
||||
|
||||
gray := gocv.NewMat()
|
||||
gocv.CvtColor(img, &gray, gocv.ColorBGRToGray)
|
||||
|
||||
cornerPoint2f := make([]gocv.Point2f, 0)
|
||||
corners := gocv.NewMatFromPoint2fVector(gocv.NewPoint2fVectorFromPoints(cornerPoint2f), true)
|
||||
|
||||
found := gocv.FindChessboardCorners(gray, boardSize, &corners,
|
||||
gocv.CalibCBAdaptiveThresh|gocv.CalibCBFastCheck|gocv.CalibCBNormalizeImage)
|
||||
|
||||
if found {
|
||||
|
||||
criteria := gocv.NewTermCriteria(gocv.Count|gocv.EPS, 30, 0.01)
|
||||
|
||||
gocv.CornerSubPix(gray, &corners, image.Pt(11, 11),
|
||||
image.Pt(-1, -1), criteria)
|
||||
|
||||
objectPoints = append(objectPoints, objp)
|
||||
imagePoints = append(imagePoints, cornerPoint2f)
|
||||
fmt.Printf("Found corners in image %d/%d\n", i+1, len(calibrationImages))
|
||||
} else {
|
||||
fmt.Printf("No corners found in image %d/%d\n", i+1, len(calibrationImages))
|
||||
}
|
||||
|
||||
img.Close()
|
||||
gray.Close()
|
||||
}
|
||||
|
||||
if len(objectPoints) < 5 {
|
||||
return nil, fmt.Errorf("need at least 5 successful calibration images, got %d", len(objectPoints))
|
||||
}
|
||||
|
||||
cameraMatrix := gocv.NewMat()
|
||||
distCoeffs := gocv.NewMat()
|
||||
rvecs := gocv.NewMat()
|
||||
tvecs := gocv.NewMat()
|
||||
|
||||
fmt.Printf("Calibrating camera with %d image pairs...\n", len(objectPoints))
|
||||
rms := gocv.CalibrateCamera(objectPointsVector, imagePointsVector, imageSize,
|
||||
&cameraMatrix, &distCoeffs, &rvecs, &tvecs, 0)
|
||||
|
||||
fmt.Printf("Camera calibration completed with RMS error: %f\n", rms)
|
||||
|
||||
profile := &CameraProfile{
|
||||
Make: "Calibrated",
|
||||
Model: "Custom",
|
||||
FocalLength: 0,
|
||||
CameraMatrix: cameraMatrix,
|
||||
DistCoeffs: distCoeffs,
|
||||
}
|
||||
|
||||
rvecs.Close()
|
||||
tvecs.Close()
|
||||
|
||||
return profile, nil
|
||||
}
|
||||
|
||||
func (pc *PhotoCorrector) BatchProcess(inputDir, outputDir string) error {
|
||||
if err := os.MkdirAll(outputDir, 0755); err != nil {
|
||||
return fmt.Errorf("failed to create output directory: %v", err)
|
||||
}
|
||||
|
||||
entries, err := os.ReadDir(inputDir)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to read input directory: %v", err)
|
||||
}
|
||||
|
||||
processedCount := 0
|
||||
errorCount := 0
|
||||
|
||||
for _, entry := range entries {
|
||||
if entry.IsDir() {
|
||||
continue
|
||||
}
|
||||
|
||||
ext := strings.ToLower(filepath.Ext(entry.Name()))
|
||||
if ext != ".jpg" && ext != ".jpeg" && ext != ".png" && ext != ".tiff" && ext != ".bmp" {
|
||||
continue
|
||||
}
|
||||
|
||||
inputPath := filepath.Join(inputDir, entry.Name())
|
||||
outputPath := filepath.Join(outputDir, entry.Name())
|
||||
|
||||
fmt.Printf("Processing: %s\n", entry.Name())
|
||||
if err := pc.ProcessImage(inputPath, outputPath); err != nil {
|
||||
log.Printf("Error processing %s: %v", entry.Name(), err)
|
||||
errorCount++
|
||||
} else {
|
||||
processedCount++
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Printf("Batch processing completed: %d successful, %d errors\n", processedCount, errorCount)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (pc *PhotoCorrector) AddCameraProfile(key string, profile CameraProfile) {
|
||||
pc.cameraProfiles[key] = profile
|
||||
}
|
||||
|
||||
func (pc *PhotoCorrector) GetSupportedCameras() []string {
|
||||
keys := make([]string, 0, len(pc.cameraProfiles))
|
||||
for k := range pc.cameraProfiles {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
return keys
|
||||
}
|
||||
|
||||
func (pc *PhotoCorrector) SaveCameraProfile(profile *CameraProfile, filename string) error {
|
||||
return fmt.Errorf("not implemented - would save profile to %s", filename)
|
||||
}
|
||||
|
||||
func (pc *PhotoCorrector) LoadCameraProfile(filename string) (*CameraProfile, error) {
|
||||
return nil, fmt.Errorf("not implemented - would load profile from %s", filename)
|
||||
}
|
||||
|
||||
func main() {
|
||||
if len(os.Args) < 2 {
|
||||
fmt.Println("Real Estate Photo Corrector")
|
||||
fmt.Println("Usage:")
|
||||
fmt.Println(" Single image: go run main.go <input_image> [output_image]")
|
||||
fmt.Println(" Batch process: go run main.go -batch <input_dir> <output_dir>")
|
||||
fmt.Println(" Show supported cameras: go run main.go -cameras")
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
pc := NewPhotoCorrector()
|
||||
defer func() {
|
||||
pc.defaultProfile.CameraMatrix.Close()
|
||||
pc.defaultProfile.DistCoeffs.Close()
|
||||
|
||||
for _, profile := range pc.cameraProfiles {
|
||||
profile.CameraMatrix.Close()
|
||||
profile.DistCoeffs.Close()
|
||||
}
|
||||
}()
|
||||
|
||||
switch os.Args[1] {
|
||||
case "-batch":
|
||||
if len(os.Args) < 4 {
|
||||
fmt.Println("Batch processing requires input and output directories")
|
||||
fmt.Println("Usage: go run main.go -batch <input_dir> <output_dir>")
|
||||
os.Exit(1)
|
||||
}
|
||||
inputDir := os.Args[2]
|
||||
outputDir := os.Args[3]
|
||||
|
||||
fmt.Printf("Batch processing images from %s to %s\n", inputDir, outputDir)
|
||||
if err := pc.BatchProcess(inputDir, outputDir); err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
|
||||
case "-cameras":
|
||||
fmt.Println("Supported camera profiles:")
|
||||
for _, camera := range pc.GetSupportedCameras() {
|
||||
fmt.Printf(" - %s\n", camera)
|
||||
}
|
||||
fmt.Printf(" - %s (default)\n", "Generic_WideAngle")
|
||||
|
||||
default:
|
||||
|
||||
inputPath := os.Args[1]
|
||||
outputPath := "corrected_" + filepath.Base(inputPath)
|
||||
if len(os.Args) > 2 {
|
||||
outputPath = os.Args[2]
|
||||
}
|
||||
|
||||
if _, err := os.Stat(inputPath); os.IsNotExist(err) {
|
||||
fmt.Printf("Error: Input file %s does not exist\n", inputPath)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
fmt.Printf("Processing single image: %s -> %s\n", inputPath, outputPath)
|
||||
fmt.Printf("Supported cameras: %v\n", pc.GetSupportedCameras())
|
||||
|
||||
if err := pc.ProcessImage(inputPath, outputPath); err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
fmt.Println("Image processing completed!")
|
||||
}
|
||||
}
|
||||
8
go.mod
Normal file
8
go.mod
Normal file
@ -0,0 +1,8 @@
|
||||
module unslop-cv
|
||||
|
||||
go 1.24.4
|
||||
|
||||
require (
|
||||
github.com/rwcarlsen/goexif v0.0.0-20190401172101-9e8deecbddbd
|
||||
gocv.io/x/gocv v0.41.0
|
||||
)
|
||||
4
go.sum
Normal file
4
go.sum
Normal file
@ -0,0 +1,4 @@
|
||||
github.com/rwcarlsen/goexif v0.0.0-20190401172101-9e8deecbddbd h1:CmH9+J6ZSsIjUK3dcGsnCnO41eRBOnY12zwkn5qVwgc=
|
||||
github.com/rwcarlsen/goexif v0.0.0-20190401172101-9e8deecbddbd/go.mod h1:hPqNNc0+uJM6H+SuU8sEs5K5IQeKccPqeSjfgcKGgPk=
|
||||
gocv.io/x/gocv v0.41.0 h1:KM+zRXUP28b6dHfhy+4JxDODbCNQNtLg8kio+YE7TqA=
|
||||
gocv.io/x/gocv v0.41.0/go.mod h1:zYdWMj29WAEznM3Y8NsU3A0TRq/wR/cy75jeUypThqU=
|
||||
Loading…
x
Reference in New Issue
Block a user