talksik
106 lines
2.6 KiB
Go
106 lines
2.6 KiB
Go
package main
|
|
|
|
import (
|
|
"fmt"
|
|
"os"
|
|
"os/exec"
|
|
"path/filepath"
|
|
"time"
|
|
)
|
|
|
|
func main() {
|
|
// Create output directory if it doesn't exist
|
|
outputDir := "output"
|
|
if err := os.MkdirAll(outputDir, 0755); err != nil {
|
|
fmt.Printf("Error creating output directory: %v\n", err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
// Generate filenames with timestamp
|
|
timestamp := time.Now().Format("20060102_150405")
|
|
originalPhoto := filepath.Join(outputDir, fmt.Sprintf("photo_%s.jpg", timestamp))
|
|
processedPhoto := filepath.Join(outputDir, fmt.Sprintf("photo_%s_nobg.png", timestamp))
|
|
|
|
fmt.Println("📸 Capturing photo from webcam...")
|
|
if err := capturePhoto(originalPhoto); err != nil {
|
|
fmt.Printf("Error capturing photo: %v\n", err)
|
|
os.Exit(1)
|
|
}
|
|
fmt.Printf("✓ Photo saved to: %s\n", originalPhoto)
|
|
|
|
fmt.Println("🎨 Removing background...")
|
|
if err := removeBackground(originalPhoto, processedPhoto); err != nil {
|
|
fmt.Printf("Error removing background: %v\n", err)
|
|
os.Exit(1)
|
|
}
|
|
fmt.Printf("✓ Processed photo saved to: %s\n", processedPhoto)
|
|
|
|
fmt.Println("\n🎉 Done! Background removed successfully!")
|
|
}
|
|
|
|
// capturePhoto captures a single frame from the webcam using ffmpeg
|
|
func capturePhoto(outputPath string) error {
|
|
// For Linux, typically /dev/video0
|
|
// For macOS, use "0:0" with avfoundation
|
|
// For Windows, use "video=CAMERA_NAME" with dshow
|
|
|
|
var cmd *exec.Cmd
|
|
|
|
// Detect platform and use appropriate video device
|
|
if _, err := os.Stat("/dev/video0"); err == nil {
|
|
// Linux
|
|
cmd = exec.Command("ffmpeg",
|
|
"-f", "v4l2",
|
|
"-i", "/dev/video0",
|
|
"-frames:v", "1",
|
|
"-y",
|
|
outputPath,
|
|
)
|
|
} else {
|
|
// Try macOS
|
|
cmd = exec.Command("ffmpeg",
|
|
"-f", "avfoundation",
|
|
"-i", "0",
|
|
"-frames:v", "1",
|
|
"-y",
|
|
outputPath,
|
|
)
|
|
}
|
|
|
|
output, err := cmd.CombinedOutput()
|
|
if err != nil {
|
|
return fmt.Errorf("ffmpeg error: %v\nOutput: %s", err, string(output))
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// removeBackground removes the background using rembg Docker container
|
|
func removeBackground(inputPath, outputPath string) error {
|
|
// Get absolute path of current directory
|
|
cwd, err := os.Getwd()
|
|
if err != nil {
|
|
return fmt.Errorf("failed to get working directory: %v", err)
|
|
}
|
|
|
|
// Docker paths are relative to /data mount point
|
|
dockerInput := "/data/" + inputPath
|
|
dockerOutput := "/data/" + outputPath
|
|
|
|
// Use rembg Docker container
|
|
cmd := exec.Command("docker", "run", "--rm",
|
|
"-v", fmt.Sprintf("%s:/data", cwd),
|
|
"danielgatis/rembg",
|
|
"i",
|
|
dockerInput,
|
|
dockerOutput,
|
|
)
|
|
|
|
output, err := cmd.CombinedOutput()
|
|
if err != nil {
|
|
return fmt.Errorf("docker/rembg error: %v\nOutput: %s", err, string(output))
|
|
}
|
|
|
|
return nil
|
|
}
|