MynooMuz commited on
Commit
ee38b6b
·
1 Parent(s): caa84f3

refactor: use official pixelfixer.api.process to match retrodiffusion behavior exactly

Browse files
Files changed (1) hide show
  1. app.py +39 -30
app.py CHANGED
@@ -18,67 +18,76 @@ def process_image(input_image_path, color_limit, export_scale):
18
  if input_image_path is None:
19
  return None, None, "Please upload an image."
20
 
21
- # Parse color limit
22
  k_colors = 0
23
  if color_limit != "Unlimited":
24
  k_colors = int(color_limit)
25
 
26
- # Parse scale multiplier
27
  scale_mult = int(export_scale.replace("x", ""))
28
 
29
- # Get original filename
30
  original_name = os.path.basename(input_image_path)
31
  base, _ = os.path.splitext(original_name)
32
-
33
- # We want to avoid generic names like 'image' if possible
34
  if base.startswith("image") and len(base) > 5:
35
  base = "image"
36
 
37
- input_image = Image.open(input_image_path)
38
- rgba = np.array(input_image.convert("RGBA"))
39
 
40
- t0 = time.time()
 
 
 
41
  try:
42
- r = detect(rgba)
 
 
 
 
 
 
 
43
  except Exception as e:
44
- return None, None, f"Error detecting pixel grid: {str(e)}"
45
 
46
- dt = round(time.time() - t0, 2)
47
-
48
- # Create a unique temp directory for this generation
49
  tmp_dir = tempfile.mkdtemp()
50
 
51
- # Generate and save overlay
52
- overlay_img = Image.fromarray(_overlay(rgba, r))
53
- overlay_path = os.path.join(tmp_dir, f"{base}-overlay.png")
54
- overlay_img.save(overlay_path, format="PNG")
55
 
56
- # Generate and save extraction with accurate colors
57
- from pixelfixer.reconstruct import two_stage_pack
58
- low = two_stage_pack(rgba, int(r["cols"]), int(r["rows"]), k_colors=k_colors)
59
-
60
- # Strictly enforce color palette if limit is set
61
  if k_colors > 0:
62
  from pixelfixer.quantize import kmeans_quantize
63
- low, _, _ = kmeans_quantize(low, k=k_colors)
64
 
65
- extract_img = Image.fromarray(low)
66
 
67
  # Upscale using perfect multiples to maintain aspect ratio and crispness
68
- target_w = int(r["cols"]) * scale_mult
69
- target_h = int(r["rows"]) * scale_mult
70
  extract_img = extract_img.resize((target_w, target_h), Image.NEAREST)
71
 
72
  extract_path = os.path.join(tmp_dir, f"{base}-pixelated.png")
73
  extract_img.save(extract_path, format="PNG")
74
 
75
- # Create info string
76
- info = f"Grid detected: {r['cols']}x{r['rows']}\n"
 
 
 
 
 
 
 
 
 
 
77
  info += f"Output size: {target_w}x{target_h} ({export_scale} scale)\n"
78
- info += f"Step size: {r['step_x']}x{r['step_y']}\n"
 
79
  if k_colors > 0:
80
  info += f"Colors: {k_colors}\n"
81
- info += f"Processing time: {dt}s"
82
 
83
  return extract_path, info, overlay_path
84
 
 
18
  if input_image_path is None:
19
  return None, None, "Please upload an image."
20
 
21
+ # Parse settings
22
  k_colors = 0
23
  if color_limit != "Unlimited":
24
  k_colors = int(color_limit)
25
 
 
26
  scale_mult = int(export_scale.replace("x", ""))
27
 
 
28
  original_name = os.path.basename(input_image_path)
29
  base, _ = os.path.splitext(original_name)
 
 
30
  if base.startswith("image") and len(base) > 5:
31
  base = "image"
32
 
33
+ # Use the official API to match the web app's exact behavior
34
+ from pixelfixer.api import process
35
 
36
+ # Read the image file as bytes
37
+ with open(input_image_path, "rb") as f:
38
+ img_bytes = f.read()
39
+
40
  try:
41
+ # Call the exact same pipeline used by the official server
42
+ result = process(
43
+ img_bytes,
44
+ mode="full",
45
+ two_stage=True,
46
+ k_colors=k_colors,
47
+ return_png=False
48
+ )
49
  except Exception as e:
50
+ return None, f"Error processing image: {str(e)}", None
51
 
52
+ # Create a unique temp directory
 
 
53
  tmp_dir = tempfile.mkdtemp()
54
 
55
+ # The result contains the pure pixel art array (1x scale)
56
+ low_res_array = result["array"]
 
 
57
 
58
+ # If the user requested a strict color limit, we still need to enforce it via k-means
 
 
 
 
59
  if k_colors > 0:
60
  from pixelfixer.quantize import kmeans_quantize
61
+ low_res_array, _, _ = kmeans_quantize(low_res_array, k=k_colors)
62
 
63
+ extract_img = Image.fromarray(low_res_array)
64
 
65
  # Upscale using perfect multiples to maintain aspect ratio and crispness
66
+ target_w = result["cols"] * scale_mult
67
+ target_h = result["rows"] * scale_mult
68
  extract_img = extract_img.resize((target_w, target_h), Image.NEAREST)
69
 
70
  extract_path = os.path.join(tmp_dir, f"{base}-pixelated.png")
71
  extract_img.save(extract_path, format="PNG")
72
 
73
+ # Optional: We can still generate the overlay using the underlying image and result data
74
+ try:
75
+ input_img = Image.open(input_image_path).convert("RGBA")
76
+ rgba = np.array(input_img)
77
+ overlay_img = Image.fromarray(_overlay(rgba, result))
78
+ overlay_path = os.path.join(tmp_dir, f"{base}-overlay.png")
79
+ overlay_img.save(overlay_path, format="PNG")
80
+ except Exception:
81
+ overlay_path = None
82
+
83
+ # Create info string exactly as reported by the official engine
84
+ info = f"Grid detected: {result['cols']}x{result['rows']}\n"
85
  info += f"Output size: {target_w}x{target_h} ({export_scale} scale)\n"
86
+ info += f"Step size: {result['step_x']}x{result['step_y']}\n"
87
+ info += f"Confidence: {result['confidence'].upper()} ({result['consensus']})\n"
88
  if k_colors > 0:
89
  info += f"Colors: {k_colors}\n"
90
+ info += f"Detect: {result['detect_s']}s | Recon: {result['recon_s']}s"
91
 
92
  return extract_path, info, overlay_path
93