Update app.py

app.py

@@ -228,14 +228,13 @@ class WarehouseStitcher:
         result_rgb = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
         result_pil = Image.fromarray(result_rgb)
         
-        # Save for download
-        buf = io.BytesIO()
-        result_pil.save(buf, format='PNG', optimize=True)
-        buf.seek(0)
-        
-        # Return (image preview, logs, downloadable file)
-        # return result_pil, "\n".join(logs), ("stitched_result.png", buf.getvalue())
-        return result_pil, "\n".join(logs), buf.getvalue()
+        # Save to temporary file for download
+        with tempfile.NamedTemporaryFile(mode='wb', suffix='.png', delete=False) as f:
+            result_pil.save(f, format='PNG', optimize=True)
+            temp_path = f.name
+        
+        # Return (image preview, logs, downloadable file path)
+        return result_pil, "\n".join(logs), temp_path
 
 
 class PoseGuidedWarehouseStitcher(WarehouseStitcher):
@@ -400,13 +399,12 @@ class PoseGuidedWarehouseStitcher(WarehouseStitcher):
         result_rgb = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
         result_pil = Image.fromarray(result_rgb)
         
-        # Save for download
-        buf = io.BytesIO()
-        result_pil.save(buf, format='PNG', optimize=True)
-        buf.seek(0)
+        # Save to temporary file for download
+        with tempfile.NamedTemporaryFile(mode='wb', suffix='.png', delete=False) as f:
+            result_pil.save(f, format='PNG', optimize=True)
+            temp_path = f.name
         
-        # return result_pil, "\n".join(logs), buf
-        return result_pil, "\n".join(logs), buf.getvalue()
+        return result_pil, "\n".join(logs), temp_path
 
 
 def create_demo():
@@ -489,21 +487,54 @@ def create_demo():
                 if len(json_files) == 0:
                     return None, "❌ No JSON metadata files found in ZIP", None
                 
-                # Match images with metadata by filename
-                image_paths = []
-                metadata_paths = []
+                # NEW: Match images with metadata by extracting base timestamp
+                # Pattern: z-9000_cam0_20251030_152626_496053_cam0.png -> z-9000_cam0_20251030_152626_496053
+                image_metadata_pairs = []
                 
                 for img_file in image_files:
-                    # Look for matching JSON
-                    json_name = img_file.stem + '.json'
+                    img_name = img_file.stem  # Filename without extension
+                    
+                    # Try to extract base name (remove _camX suffix if present)
+                    # Pattern: z-9000_cam0_20251030_152626_496053_cam0 -> z-9000_cam0_20251030_152626_496053
+                    if '_cam' in img_name:
+                        parts = img_name.rsplit('_cam', 1)
+                        if len(parts) == 2:
+                            base_name = parts[0]  # Everything before last _camX
+                        else:
+                            base_name = img_name
+                    else:
+                        base_name = img_name
+                    
+                    # Look for matching JSON with base name
+                    json_name = base_name + '.json'
                     json_candidates = [j for j in json_files if j.name == json_name]
                     
                     if json_candidates:
-                        image_paths.append(str(img_file))
-                        metadata_paths.append(str(json_candidates[0]))
+                        image_metadata_pairs.append({
+                            'image': str(img_file),
+                            'metadata': str(json_candidates[0]),
+                            'base_name': base_name,
+                            'timestamp': None  # Will be loaded from JSON
+                        })
+                
+                if len(image_metadata_pairs) < 2:
+                    return None, f"❌ Only {len(image_metadata_pairs)} images have matching metadata. Need at least 2.\n\nFound images: {[f.name for f in image_files[:5]]}\nFound JSON: {[j.name for j in json_files[:5]]}", None
+                
+                # Load timestamps and sort
+                for pair in image_metadata_pairs:
+                    try:
+                        with open(pair['metadata'], 'r') as f:
+                            metadata = json.load(f)
+                            pair['timestamp'] = metadata['nav_snapshot']['timestamp_usec']
+                    except Exception as e:
+                        return None, f"❌ Error reading metadata {pair['metadata']}: {str(e)}", None
+                
+                # Sort by timestamp
+                image_metadata_pairs.sort(key=lambda x: x['timestamp'])
                 
-                if len(image_paths) < 2:
-                    return None, f"❌ Only {len(image_paths)} images have matching metadata", None
+                # Extract sorted lists
+                image_paths = [p['image'] for p in image_metadata_pairs]
+                metadata_paths = [p['metadata'] for p in image_metadata_pairs]
                 
                 # Update configuration
                 pose_stitcher.config['feature_extractor'] = feature_type
@@ -516,7 +547,8 @@ def create_demo():
                 return pose_stitcher.stitch_with_poses(image_paths, metadata_paths)
                 
         except Exception as e:
-            return None, f"❌ Error processing ZIP: {str(e)}", None
+            import traceback
+            return None, f"❌ Error processing ZIP: {str(e)}\n\n{traceback.format_exc()}", None
     
     # Custom CSS
     custom_css = """
@@ -695,8 +727,7 @@ dataset.zip
                 )
                 
                 download_btn = gr.File(
-                    label="⬇️ Download High-Resolution Result",
-                    type="binary"
+                    label="⬇️ Download High-Resolution Result"
                 )
                 
             with gr.Column(scale=1):
@@ -742,6 +773,6 @@ if __name__ == "__main__":
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,
-        share=False,
+        share=True,
         show_error=True
     )