commit 7b6028330339f43288c673b42c9648e0127a20f5 Author: Mark Bailey Date: Mon Jul 14 00:08:49 2025 -0400 Initial commit diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..7e03544 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +unslop-cv diff --git a/cmd/unslop-cv.go b/cmd/unslop-cv.go new file mode 100644 index 0000000..3672a7a --- /dev/null +++ b/cmd/unslop-cv.go @@ -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 [output_image]") + fmt.Println(" Batch process: go run main.go -batch ") + 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 ") + 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!") + } +} diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..87e725f --- /dev/null +++ b/go.mod @@ -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 +) diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..0cf44d3 --- /dev/null +++ b/go.sum @@ -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=