Updated app.py, possible fix for coding module integration

app.py

@@ -5,38 +5,32 @@ import traceback
 
 import gradio as gr
 from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse, FileResponse
 
 from visualization import process_wireframe
 from CodingModule import HTMLGenerator
 
-
-# ============================================================
-# DIRECTORIES
-# ============================================================
-
-TEMP_DIR = "./temp"
-OUTPUT_DIR = "./output"
-
-os.makedirs(TEMP_DIR, exist_ok=True)
-os.makedirs(OUTPUT_DIR, exist_ok=True)
-
-
-# ============================================================
-# FASTAPI BACKEND
-# ============================================================
-
+# -----------------------------------------------------------------------------
+# FASTAPI (for Firebase / programmatic access)
+# -----------------------------------------------------------------------------
 api = FastAPI()
 
 api.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # change to your Firebase domain later
+    allow_origins=["*"], #Replace with the Firebase domain
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
+TEMP_DIR = "./temp"
+OUTPUT_DIR = "./output"
+
+os.makedirs(TEMP_DIR, exist_ok=True)
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+
 
 @api.get("/")
 def health_check():
@@ -49,24 +43,9 @@ async def process_wireframe_api(image: UploadFile = File(...)):
     temp_path = os.path.join(TEMP_DIR, f"{file_id}_{image.filename}")
 
     try:
-        # ----------------------------------------------------
-        # SAVE UPLOADED IMAGE
-        # ----------------------------------------------------
         with open(temp_path, "wb") as f:
             shutil.copyfileobj(image.file, f)
 
-        file_size = os.path.getsize(temp_path)
-        print("Uploaded file size:", file_size)
-
-        if file_size == 0:
-            return JSONResponse(
-                status_code=400,
-                content={"error": "Uploaded file is empty — check frontend upload format."}
-            )
-
-        # ----------------------------------------------------
-        # RUN DETECTION PIPELINE
-        # ----------------------------------------------------
         results = process_wireframe(
             image_path=temp_path,
             save_json=True,
@@ -74,31 +53,26 @@ async def process_wireframe_api(image: UploadFile = File(...)):
             show_visualization=False
         )
 
-        if not results or not results.get("normalized_elements"):
+        if not results or not results.get('normalized_elements'):
             return JSONResponse(
                 status_code=400,
                 content={"error": "No elements detected"}
             )
-
+        
         json_path = results.get("json_path")
-
-        # ----------------------------------------------------
-        # GENERATE STYLED HTML
-        # ----------------------------------------------------
+        
+        #Generate styled HTML using CodingModule
         try:
+            html_generator = HTMLGenerator(json_path)
             html_filename = f"{file_id}_styled.html"
             html_path = os.path.join(OUTPUT_DIR, html_filename)
-
-            generator = HTMLGenerator(json_path)
-            generator.generate_html(html_path)
-
+            html_generator.generate_html(html_path)
         except Exception as e:
-            print("HTML generation error:", e)
+            print(f"Warning: Could not generate styled HTML: {e}")
             html_path = results.get("html_path")
-
+        
         return {
             "success": True,
-            "file_id": file_id,
             "json_path": json_path,
             "html_path": html_path,
             "total_elements": len(results["normalized_elements"])
@@ -115,97 +89,155 @@ async def process_wireframe_api(image: UploadFile = File(...)):
         if os.path.exists(temp_path):
             os.remove(temp_path)
 
-
-# ============================================================
-# SERVE GENERATED FILES (HTML / JSON)
-# ============================================================
-
-@api.get("/output/{filename}")
-async def serve_output(filename: str):
-    path = os.path.join(OUTPUT_DIR, filename)
-    if not os.path.exists(path):
-        return JSONResponse(status_code=404,
-                            content={"error": "File not found"})
-
-    if filename.endswith(".html"):
-        return FileResponse(path, media_type="text/html")
-
-    if filename.endswith(".json"):
-        return FileResponse(path, media_type="application/json")
-
-    return FileResponse(path)
-
-
-# ============================================================
-# GRADIO FRONTEND (for HF testing only)
-# ============================================================
-
+#Endpoint to serve generated HTML files
+@api.get("/view-html/{file_id}")
+async def serve_output_file(filename: str):
+    file_path = os.path.join(OUTPUT_DIR, filename)
+    if os.path.exists(file_path):
+        if filename.endswith('.html'):
+            return FileResponse(file_path, media_type="text/html")
+        elif filename.endswith('.json'):
+            return FileResponse(file_path, media_type="application/json")
+        else:
+            return FileResponse(file_path)
+    return JSONResponse(status_code=404, content={"error": "File not Found"})
+
+# # Legacy endpoint for just for backward compatibility
+# @api.get("/view-html/{file_id}")
+# async def view_html(file_id: str):
+#     html_path = os.path.join(OUTPUT_DIR, f"{file_id}.html")
+#     if os.path.exists(html_path):
+#         return FileResponse(html_path, media_type="text/html")
+#     # Try with _styled suffix
+#     html_path = os.path.join(OUTPUT_DIR, f"{file_id}_styled.html")
+#     if os.path.exists(html_path):
+#         return FileResponse(html_path, media_type="text/html")
+#     return JSONResponse(status_code=404, content={"error": "File not found"})
+
+# -----------------------------------------------------------------------------
+# GRADIO (for Hugging Face UI)
+# -----------------------------------------------------------------------------
 def gradio_process(image):
-
+    """
+    Gradio passes a PIL Image.
+    We save it temporarily and reuse the SAME pipeline.
+    """
     if image is None:
-        return "Upload an image first", None, None
-
+        return "Please upload an image", None, None
+    
     file_id = str(uuid.uuid4())
     temp_path = os.path.join(TEMP_DIR, f"{file_id}.png")
-
-    image.save(temp_path)
+    
+    try: 
+        image.save(temp_path)
+    except Exception as e:
+        return f"Error saving image: {str(e)}", None, None
 
     try:
         results = process_wireframe(
             image_path=temp_path,
             save_json=True,
-            save_html=True,
+            save_html=False,
             show_visualization=False
         )
 
-        if not results or not results.get("normalized_elements"):
+        if not results or not results.get('normalized_elements'):
             return "No elements detected", None, None
 
         json_path = results.get("json_path")
-
+        
+        #Generate styled HTML using CodingModule
         try:
+            html_generator = HTMLGenerator(json_path)
             html_filename = f"{file_id}_styled.html"
             html_path = os.path.join(OUTPUT_DIR, html_filename)
-
-            generator = HTMLGenerator(json_path)
-            generator.generate_html(html_path)
-
-        except Exception:
+            html_generator.generate_html(html_path)
+        except Exception as e:
+            print(f"Warning: Could not generate styled HTML: {e}")
+            traceback.print_exc()
             html_path = results.get("html_path")
-
-        return (
-            f"Detected {len(results['normalized_elements'])} elements",
-            json_path,
-            html_path
-        )
+        
+        num_elements = len(results['normalized_elements'])
+        status_msg = f"Successfully detected {num_elements} elemnts\n"
+        status_msg += f"JSON: {os.path.basename(json_path)}\n"
+        status_msg += f"HTML: {os.path.basename(html_path)}"
+        
+        return status_msg, json_path, html_path
 
     except Exception as e:
         traceback.print_exc()
-        return f"Error: {e}", None, None
+        return f"Error: {str(e)}", None, None
 
     finally:
-        os.remove(temp_path)
-
+        if os.path.exists(temp_path):
+            os.remove(temp_path)
 
-with gr.Blocks(title="Wireframe Layout Normalizer") as demo:
+# Enhanced Gradio interface with better UI
+with gr.Blocks(theme=gr.themes.Soft(), title="Wireframe Layout Normalizer") as demo:
     gr.Markdown("# Wireframe Layout Normalizer")
-
-    inp = gr.Image(type="pil")
-    btn = gr.Button("Process")
-    out1 = gr.Textbox(label="Status")
-    out2 = gr.File(label="JSON")
-    out3 = gr.File(label="HTML")
-
-    btn.click(gradio_process, inp, [out1, out2, out3])
-
-
-# ============================================================
-# MOUNT FASTAPI + GRADIO
-# ============================================================
-
+    gr.Markdown("Upload a wireframe image to extract and normalize UI layout elements")
+    
+    gr.Markdown("""
+    ###  Academic Research Project made by SeroTech 2025-2026 Holy Angel University
+    This tool is part of a thesis project. Upload your wireframe sketch to automatically:
+    -  Detect UI elements (buttons, text fields, images, etc.)
+    -  Normalize positions and alignments
+    -  Export structured JSON data
+    -  Generate beautiful HTML preview
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            image_input = gr.Image(type="pil", label="Upload Wireframe Image")
+            process_btn = gr.Button("Process Wireframe", variant="primary", size="lg")
+            
+            gr.Markdown("""
+            #### Tips:
+            - Upload clear wireframe sketches
+            - PNG, JPG, or JPEG formats supported
+            - Higher resolution = better detection
+            """)
+            
+        with gr.Column(scale=1):
+            status_output = gr.Textbox(
+                label="Processing Status", 
+                lines=3,
+                placeholder="Upload an image and click Process to begin..."
+            )
+            json_output = gr.File(label="Normalized JSON Output")
+            html_output = gr.File(label="Generated HTML Preview")
+    
+    process_btn.click(
+        fn=gradio_process,
+        inputs=image_input,
+        outputs=[status_output, json_output, html_output]
+    )
+    
+    gr.Markdown("""
+    ---
+    ### How to use:
+    1. *Upload* a wireframe image (hand-drawn or digital)
+    2. *Click* "Process Wireframe" and wait for processing
+    3. *Download* the JSON file for structured data
+    4. *Download* the HTML file to see a styled preview
+    
+    ### Output Files:
+    - *JSON*: Contains all detected elements with positions, types, and grid data
+    - *HTML*: Beautiful, responsive preview with modern styling
+    
+    ### Technical Details:
+    - Model: TensorFlow-based object detection
+    - Grid System: 24-column responsive layout
+    - Supported Elements: Buttons, Checkboxes, Text Fields, Images, Paragraphs, Text, Navbars
+    """)
+
+
+# -----------------------------------------------------------------------------
+# ENTRY POINT (THIS IS IMPORTANT)
+# -----------------------------------------------------------------------------
 app = gr.mount_gradio_app(api, demo, path="/")
 
-
-if __name__ == "__main__":
+if _name_ == "_main_":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    # For local testing
+    uvicorn.run(app, host="0.0.0.0", port=7860)