Add progress bar

main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"flag"
 	"fmt"
 	"image"
 	"image/color"
@@ -8,20 +9,43 @@ import (
 	_ "image/jpeg"
 	"image/png"
 	_ "image/png"
+	"log"
 	"math"
 	"os"
 	"sort"
 	"time"
 
 	"github.com/nfnt/resize"
+	"gopkg.in/cheggaaa/pb.v1"
 )
 
-const (
+// MaxImageSize is the resized image maximum width or height depending on the image ratio.
-	BlockSize         int = 4
+const MaxImageSize = 320
-	DistanceThreshold     = 0.4
+
-	OffsetThreshold       = 72
+const Banner = `
-	ForgeryThreshold      = 170
+  __                          _
-	MaxImageSize          = 320
+ / _| ___  _ __ ___ _ __  ___(_) ___
+| |_ / _ \| '__/ _ \ '_ \/ __| |/ __|
+|  _| (_) | | |  __/ | | \__ \ | (__
+|_|  \___/|_|  \___|_| |_|___/_|\___|
+
+Image forgery detection library.
+    Version: %s
+
+`
+
+// Version indicates the current build version.
+var Version string
+
+var (
+	// Flags
+	source            = flag.String("in", "", "Source")
+	destination       = flag.String("out", "", "Destination")
+	blurRadius        = flag.Int("blur", 1, "Blur radius")
+	blockSize         = flag.Int("bs", 4, "Block size")
+	offsetThreshold   = flag.Int("ot", 72, "Offset threshold")
+	distanceThreshold = flag.Float64("dt", 0.4, "Distance threshold")
+	forgeryThreshold  = flag.Float64("ft", 210, "Forgery threshold")
 )
 
 // pixel struct contains the discrete cosine transformation R,G,B,Y values.
@@ -69,9 +93,20 @@ var (
 )
 
 func main() {
+	done := make(chan struct{})
+	flag.Usage = func() {
+		fmt.Fprintf(os.Stderr, fmt.Sprintf(Banner, Version))
+		flag.PrintDefaults()
+	}
+	flag.Parse()
+
+	if len(*source) == 0 || len(*destination) == 0 {
+		log.Fatal("Usage: forensic -in input.jpg -out out.jpg")
+	}
+
 	start := time.Now()
 
-	input, err := os.Open("korea_forged.jpg")
+	input, err := os.Open(*source)
 	defer input.Close()
 
 	if err != nil {
@@ -90,12 +125,26 @@ func main() {
 		resizedImg = src
 	}
 
-	img := imgToNRGBA(resizedImg)
+	go func() {
+		result := process(resizedImg, done)
+		if result {
+			fmt.Println("\nThe image is forged!")
+		} else {
+			fmt.Println("\nThe image is not forged!")
+		}
+	}()
+	<-done
+
+	fmt.Printf("\nDone in: %.2fs\n", time.Since(start).Seconds())
+}
+
+func process(input image.Image, done chan struct{}) bool {
+	img := imgToNRGBA(input)
 	output := image.NewRGBA(img.Bounds())
 	draw.Draw(output, image.Rect(0, 0, img.Bounds().Dx(), img.Bounds().Dy()), img, image.ZP, draw.Src)
 
 	// Blur the image to eliminate the details.
-	blurImg := StackBlur(img, 1)
+	blurImg := StackBlur(img, uint32(*blurRadius))
 
 	// Convert image to YUV color space
 	yuv := convertRGBImageToYUV(blurImg)
@@ -103,19 +152,21 @@ func main() {
 	draw.Draw(newImg, image.Rect(0, 0, yuv.Bounds().Dx(), yuv.Bounds().Dy()), yuv, image.ZP, draw.Src)
 
 	dx, dy := yuv.Bounds().Max.X, yuv.Bounds().Max.Y
-	bdx, bdy := (dx - BlockSize + 1), (dy - BlockSize + 1)
+	bdx, bdy := (dx - *blockSize + 1), (dy - *blockSize + 1)
 	n := math.Max(float64(dx), float64(dy))
 
 	var blocks []imageBlock
 	for i := 0; i < bdx; i++ {
 		for j := 0; j < bdy; j++ {
-			r := image.Rect(i, j, i+BlockSize, j+BlockSize)
+			r := image.Rect(i, j, i+*blockSize, j+*blockSize)
 			block := newImg.SubImage(r).(*image.RGBA)
 			blocks = append(blocks, imageBlock{x: i, y: j, img: block})
-			//draw.Draw(newImg, image.Rect(0, 0, yuv.Bounds().Max.X, yuv.Bounds().Max.Y), block, image.ZP, draw.Src)
 		}
 	}
 
+	bar := pb.StartNew(len(blocks))
+	bar.Prefix("Generate: ")
+
 	for _, block := range blocks {
 		// Average RGB value.
 		var avr, avg, avb float64
@@ -124,23 +175,23 @@ func main() {
 		i0 := b.PixOffset(b.Bounds().Min.X, b.Bounds().Min.Y)
 		i1 := i0 + b.Bounds().Dx()*4
 
-		dctPixels := make(dctPx, BlockSize*BlockSize)
+		dctPixels := make(dctPx, *blockSize**blockSize)
-		for u := 0; u < BlockSize; u++ {
+		for u := 0; u < *blockSize; u++ {
-			dctPixels[u] = make([]pixel, BlockSize)
+			dctPixels[u] = make([]pixel, *blockSize)
-			for v := 0; v < BlockSize; v++ {
+			for v := 0; v < *blockSize; v++ {
 				for i := i0; i < i1; i += 4 {
 					// Get the YUV converted image pixels
 					yc, uc, vc, _ := b.Pix[i+0], b.Pix[i+2], b.Pix[i+2], b.Pix[i+3]
 					// Convert YUV to RGB and obtain the R value
 					r, g, b := color.YCbCrToRGB(yc, uc, vc)
 
-					for x := 0; x < BlockSize; x++ {
+					for x := 0; x < *blockSize; x++ {
-						for y := 0; y < BlockSize; y++ {
+						for y := 0; y < *blockSize; y++ {
 							// Compute Discrete Cosine coefficients
-							cr += dct(float64(x), float64(y), float64(u), float64(v), float64(BlockSize)) * float64(r)
+							cr += dct(float64(x), float64(y), float64(u), float64(v), float64(*blockSize)) * float64(r)
-							cg += dct(float64(x), float64(y), float64(u), float64(v), float64(BlockSize)) * float64(g)
+							cg += dct(float64(x), float64(y), float64(u), float64(v), float64(*blockSize)) * float64(g)
-							cb += dct(float64(x), float64(y), float64(u), float64(v), float64(BlockSize)) * float64(b)
+							cb += dct(float64(x), float64(y), float64(u), float64(v), float64(*blockSize)) * float64(b)
-							cy += dct(float64(x), float64(y), float64(u), float64(v), float64(BlockSize)) * float64(yc)
+							cy += dct(float64(x), float64(y), float64(u), float64(v), float64(*blockSize)) * float64(yc)
 
 							avr += float64(r)
 							avg += float64(g)
@@ -167,15 +218,17 @@ func main() {
 				dctPixels[u][v] = pixel{cr, cg, cb, cy}
 
 				// Get the quantized DCT coefficients.
-				dctPixels[u][v].r = (dctPixels[u][v].r / q4x4[u][v])
+				if *blockSize <= 4 {
-				dctPixels[u][v].g = (dctPixels[u][v].g / q4x4[u][v])
+					dctPixels[u][v].r = dctPixels[u][v].r / q4x4[u][v]
-				dctPixels[u][v].b = (dctPixels[u][v].b / q4x4[u][v])
+					dctPixels[u][v].g = dctPixels[u][v].g / q4x4[u][v]
-				dctPixels[u][v].y = (dctPixels[u][v].y / q4x4[u][v])
+					dctPixels[u][v].b = dctPixels[u][v].b / q4x4[u][v]
+					dctPixels[u][v].y = dctPixels[u][v].y / q4x4[u][v]
+				}
 			}
 		}
-		avr /= float64(BlockSize * BlockSize)
+		avr /= float64(*blockSize * *blockSize)
-		avg /= float64(BlockSize * BlockSize)
+		avg /= float64(*blockSize * *blockSize)
-		avb /= float64(BlockSize * BlockSize)
+		avb /= float64(*blockSize * *blockSize)
 
 		features = append(features, feature{x: block.x, y: block.y, coef: dctPixels[0][0].y})
 		features = append(features, feature{x: block.x, y: block.y, coef: dctPixels[0][1].y})
@@ -188,11 +241,16 @@ func main() {
 		features = append(features, feature{x: block.x, y: block.y, coef: avr})
 		features = append(features, feature{x: block.x, y: block.y, coef: avb})
 		features = append(features, feature{x: block.x, y: block.y, coef: avg})
+		bar.Increment()
 	}
+	bar.Finish()
 
 	// Lexicographically sort the feature vectors
 	sort.Sort(featVec(features))
 
+	bar = pb.StartNew(len(features)-1)
+	bar.Prefix("Analyze: ")
+
 	for i := 0; i < len(features)-1; i++ {
 		blockA, blockB := features[i], features[i+1]
 		result := analyzeBlocks(blockA, blockB)
@@ -200,7 +258,9 @@ func main() {
 		if result != nil {
 			vectors = append(vectors, *result)
 		}
+		bar.Increment()
 	}
+	bar.Finish()
 
 	simBlocks := getSuspiciousBlocks(vectors)
 	forgedBlocks, result := filterOutNeighbors(simBlocks)
@@ -209,13 +269,13 @@ func main() {
 	overlay := color.RGBA{255, 0, 0, 255}
 
 	for _, bl := range forgedBlocks {
-		draw.Draw(forgedImg, image.Rect(bl.xa, bl.ya, bl.xa+BlockSize*2, bl.ya+BlockSize*2), &image.Uniform{overlay}, image.ZP, draw.Over)
+		draw.Draw(forgedImg, image.Rect(bl.xa, bl.ya, bl.xa+*blockSize*2, bl.ya+*blockSize*2), &image.Uniform{overlay}, image.ZP, draw.Over)
 	}
 
 	final := StackBlur(imgToNRGBA(forgedImg), 10)
 	draw.Draw(output, img.Bounds(), final, image.ZP, draw.Over)
 
-	out, err := os.Create("output.png")
+	out, err := os.Create(*destination)
 	if err != nil {
 		fmt.Printf("Error creating output file: %v", err)
 	}
@@ -224,8 +284,8 @@ func main() {
 		fmt.Printf("Error encoding image file: %v", err)
 	}
 
-	fmt.Println("\n", result)
+	done <- struct {}{}
-	fmt.Printf("\nDone in: %.2fs\n", time.Since(start).Seconds())
+	return result
 }
 
 //convertRGBImageToYUV coverts the image from RGB to YUV color space.
@@ -260,7 +320,7 @@ func analyzeBlocks(blockA, blockB feature) *vector {
 		offsetY: math.Abs(dy),
 	}
 
-	if dist < DistanceThreshold {
+	if dist < *distanceThreshold {
 		return res
 	}
 	return nil
@@ -279,6 +339,8 @@ func getSuspiciousBlocks(vect []vector) newVector {
 	//For each pair of candidate compute the accumulative number of the corresponding shift vectors.
 	duplicates := make(map[offset]int)
 
+	bar := pb.StartNew(len(vect)).Prefix("Detect: ")
+
 	for _, v := range vect {
 		// Check for duplicate blocks
 		offsetX := v.offsetX
@@ -294,13 +356,14 @@ func getSuspiciousBlocks(vect []vector) newVector {
 
 		// If the accumulative number of corresponding shift vectors is greater than
 		// a predefined threshold, the corresponding regions are marked as suspicious.
-		if duplicates[*offset] > OffsetThreshold {
+		if duplicates[*offset] > *offsetThreshold {
 			suspiciousBlocks = append(suspiciousBlocks, vector{
 				v.xa, v.ya, v.xb, v.yb, v.offsetX, v.offsetY,
 			})
 		}
+		bar.Increment()
 	}
-	fmt.Println("Blocks: ", len(suspiciousBlocks))
+	bar.Finish()
 	return suspiciousBlocks
 }
 
@@ -309,6 +372,8 @@ func filterOutNeighbors(vect []vector) (newVector, bool) {
 	var forgedBlocks newVector
 	var isForged bool
 
+	bar := pb.StartNew(len(vect)).Prefix("Filter: ")
+
 	for i := 1; i < len(vect); i++ {
 		blockA, blockB := vect[i-1], vect[i]
 
@@ -321,7 +386,7 @@ func filterOutNeighbors(vect []vector) (newVector, bool) {
 
 			// Evaluate the euclidean distance distance between two regions
 			// and make sure the distance is greater than a predefined threshold.
-			if dist > ForgeryThreshold {
+			if dist > *forgeryThreshold {
 				forgedBlocks = append(forgedBlocks, vector{
 					blockA.xa, blockA.ya, blockA.xb, blockA.yb, blockA.offsetX, vect[i].offsetY,
 				})
@@ -331,7 +396,9 @@ func filterOutNeighbors(vect []vector) (newVector, bool) {
 				}
 			}
 		}
+		bar.Increment()
 	}
+	bar.Finish()
 	return forgedBlocks, isForged
 }