Update app.py

app.py

@@ -1,162 +1,162 @@
-
-
-from PIL import Image
-# import torch
-import os
-import gradio as gr
-
-# Below is the code refactored into a Python class for better modularity and reusability.
-
-
-
-from transformers import TextStreamer
-
-class FloorPlanAnalyzer:
-    def __init__(self, model_path, load_in_4bit=True, gradient_checkpointing="unsloth", device="cpu"):
-        """
-        Initializes the FloorPlanAnalyzer with the specified model and configuration.
-        """
-        from unsloth import FastVisionModel  # Assuming unsloth package is installed
-
-        self.device = device
-        self.model, self.tokenizer = FastVisionModel.from_pretrained(
-            model_path,
-            load_in_4bit=load_in_4bit,
-            use_gradient_checkpointing=gradient_checkpointing,
-        )
-        FastVisionModel.for_inference(self.model)
-
-    def prepare_input(self, image_path, instruction):
-        """
-        Prepares the input for the model by loading the image and applying the chat template.
-
-        Args:
-            image_path (str): Path to the floor plan image.
-            instruction (str): Instruction text to guide the analysis.
-
-        Returns:
-            torch.Tensor: Processed inputs for the model.
-        """
-        # Load image
-        image = Image.open(image_path).convert("RGB")
-
-        # Create message template
-        messages = [
-            {"role": "user", "content": [
-                {"type": "image"},
-                {"type": "text", "text": instruction}
-            ]}
-        ]
-
-        # Generate input text
-        input_text = self.tokenizer.apply_chat_template(messages, add_generation_prompt=True)
-
-        # Tokenize and prepare inputs
-        inputs = self.tokenizer(
-            image,
-            input_text,
-            add_special_tokens=False,
-            return_tensors="pt",
-        ).to(self.device)
-
-        return inputs
-
-    def analyze(self, image_path, instruction, max_new_tokens=512, temperature=1.5, min_p=0.1):
-        """
-        Analyzes the floor plan based on the provided instruction.
-
-        Args:
-            image_path (str): Path to the floor plan image.
-            instruction (str): Instruction guiding the analysis.
-            max_new_tokens (int): Maximum number of tokens to generate.
-            temperature (float): Sampling temperature for generation.
-            min_p (float): Minimum probability for nucleus sampling.
-
-        Returns:
-            str: The generated output from the model.
-        """
-        # Prepare inputs
-        inputs = self.prepare_input(image_path, instruction)
-
-        # Set up text streamer
-        text_streamer = TextStreamer(self.tokenizer, skip_prompt=True)
-
-        # Generate output
-        output = self.model.generate(
-            **inputs,
-            streamer=text_streamer,
-            max_new_tokens=max_new_tokens,
-            use_cache=True,
-            temperature=temperature,
-            min_p=min_p,
-        )
-
-        return output
-
-
-
-
-# Instantiate the FloorPlanAnalyzer
-model_path = "./model/"
-analyzer = FloorPlanAnalyzer(model_path=model_path)
-
-
-# Sample images for Gradio examples
-# Define sample images for user convenience
-sample_images = [
-    "./sample/10_2.jpg",
-    "./sample/10_10.jpg",
-    "./sample/0_10.jpg",
-    "./sample/2_12.jpg"
-]
-
-# Ensure sample images directory exists
-os.makedirs("samples", exist_ok=True)
-# Save some dummy sample images if they don't exist (you should replace these with actual images)
-for i, sample in enumerate(sample_images):
-    if not os.path.exists(sample):
-        img = Image.new("RGB", (224, 224), color=(i * 50, i * 50, i * 50))
-        img.save(sample)
-
-# Gradio prediction function
-def predict_image(image, instruction):
-    """
-    Processes the uploaded image and instruction through the FloorPlanAnalyzer.
-
-    Args:
-        image (PIL.Image.Image): The uploaded floor plan image.
-        instruction (str): The user-provided instruction.
-
-    Returns:
-        str: The generated output description.
-    """
-    return analyzer.analyze(image, instruction)
-
-gr_interface = gr.Interface(
-    fn=predict_image,
-    inputs=[
-        gr.Image(type="pil", label="Upload Floor Plan Image"),
-        gr.Textbox(
-            label="Instruction Text",
-            value="""You are an expert in architecture and interior design. Analyze the floor plan image and describe accurately the key features, room count, layout, and any other important details you observe."""
-        )
-    ],
-    outputs=gr.Textbox(label="Analysis Result"),
-    title="Floor Plan Analyzer",
-    description="Upload a floor plan image and provide instructions to analyze it. Get detailed insights into the layout and design.",
-    examples=sample_images  # Add the examples here
-)
-
-# # Gradio UI setup with examples
-# gr_interface = gr.Interface(
-#     fn=predict_image, 
-#     inputs=gr.Image(type="pil"),  # Updated to gr.Image for image input
-#     outputs=[gr.Image(type="pil"), gr.Textbox()],  # Updated to gr.Image and gr.Textbox
-#     title="House CAD Design Object Detection", 
-#     description="Upload a CAD design image of a house to detect objects with bounding boxes and probabilities.",
-#     examples=sample_images  # Add the examples here
-# )
-
-# Launch the Gradio interface if run as main
-if __name__ == "__main__":
-    gr_interface.launch()
+
+
+from PIL import Image
+# import torch
+import os
+import gradio as gr
+
+# Below is the code refactored into a Python class for better modularity and reusability.
+
+
+import torch
+from transformers import TextStreamer
+
+class FloorPlanAnalyzer:
+    def __init__(self, model_path, load_in_4bit=True, gradient_checkpointing="unsloth", device="cpu"):
+        """
+        Initializes the FloorPlanAnalyzer with the specified model and configuration.
+        """
+        from unsloth import FastVisionModel  # Assuming unsloth package is installed
+
+        self.device = device
+        self.model, self.tokenizer = FastVisionModel.from_pretrained(
+            model_path,
+            load_in_4bit=load_in_4bit,
+            use_gradient_checkpointing=gradient_checkpointing,
+        )
+        FastVisionModel.for_inference(self.model)
+
+    def prepare_input(self, image_path, instruction):
+        """
+        Prepares the input for the model by loading the image and applying the chat template.
+
+        Args:
+            image_path (str): Path to the floor plan image.
+            instruction (str): Instruction text to guide the analysis.
+
+        Returns:
+            torch.Tensor: Processed inputs for the model.
+        """
+        # Load image
+        image = Image.open(image_path).convert("RGB")
+
+        # Create message template
+        messages = [
+            {"role": "user", "content": [
+                {"type": "image"},
+                {"type": "text", "text": instruction}
+            ]}
+        ]
+
+        # Generate input text
+        input_text = self.tokenizer.apply_chat_template(messages, add_generation_prompt=True)
+
+        # Tokenize and prepare inputs
+        inputs = self.tokenizer(
+            image,
+            input_text,
+            add_special_tokens=False,
+            return_tensors="pt",
+        ).to(self.device)
+
+        return inputs
+
+    def analyze(self, image_path, instruction, max_new_tokens=512, temperature=1.5, min_p=0.1):
+        """
+        Analyzes the floor plan based on the provided instruction.
+
+        Args:
+            image_path (str): Path to the floor plan image.
+            instruction (str): Instruction guiding the analysis.
+            max_new_tokens (int): Maximum number of tokens to generate.
+            temperature (float): Sampling temperature for generation.
+            min_p (float): Minimum probability for nucleus sampling.
+
+        Returns:
+            str: The generated output from the model.
+        """
+        # Prepare inputs
+        inputs = self.prepare_input(image_path, instruction)
+
+        # Set up text streamer
+        text_streamer = TextStreamer(self.tokenizer, skip_prompt=True)
+
+        # Generate output
+        output = self.model.generate(
+            **inputs,
+            streamer=text_streamer,
+            max_new_tokens=max_new_tokens,
+            use_cache=True,
+            temperature=temperature,
+            min_p=min_p,
+        )
+
+        return output
+
+
+
+
+# Instantiate the FloorPlanAnalyzer
+model_path = "./model/"
+analyzer = FloorPlanAnalyzer(model_path=model_path)
+
+
+# Sample images for Gradio examples
+# Define sample images for user convenience
+sample_images = [
+    "./sample/10_2.jpg",
+    "./sample/10_10.jpg",
+    "./sample/0_10.jpg",
+    "./sample/2_12.jpg"
+]
+
+# Ensure sample images directory exists
+os.makedirs("samples", exist_ok=True)
+# Save some dummy sample images if they don't exist (you should replace these with actual images)
+for i, sample in enumerate(sample_images):
+    if not os.path.exists(sample):
+        img = Image.new("RGB", (224, 224), color=(i * 50, i * 50, i * 50))
+        img.save(sample)
+
+# Gradio prediction function
+def predict_image(image, instruction):
+    """
+    Processes the uploaded image and instruction through the FloorPlanAnalyzer.
+
+    Args:
+        image (PIL.Image.Image): The uploaded floor plan image.
+        instruction (str): The user-provided instruction.
+
+    Returns:
+        str: The generated output description.
+    """
+    return analyzer.analyze(image, instruction)
+
+gr_interface = gr.Interface(
+    fn=predict_image,
+    inputs=[
+        gr.Image(type="pil", label="Upload Floor Plan Image"),
+        gr.Textbox(
+            label="Instruction Text",
+            value="""You are an expert in architecture and interior design. Analyze the floor plan image and describe accurately the key features, room count, layout, and any other important details you observe."""
+        )
+    ],
+    outputs=gr.Textbox(label="Analysis Result"),
+    title="Floor Plan Analyzer",
+    description="Upload a floor plan image and provide instructions to analyze it. Get detailed insights into the layout and design.",
+    examples=sample_images  # Add the examples here
+)
+
+# # Gradio UI setup with examples
+# gr_interface = gr.Interface(
+#     fn=predict_image, 
+#     inputs=gr.Image(type="pil"),  # Updated to gr.Image for image input
+#     outputs=[gr.Image(type="pil"), gr.Textbox()],  # Updated to gr.Image and gr.Textbox
+#     title="House CAD Design Object Detection", 
+#     description="Upload a CAD design image of a house to detect objects with bounding boxes and probabilities.",
+#     examples=sample_images  # Add the examples here
+# )
+
+# Launch the Gradio interface if run as main
+if __name__ == "__main__":
+    gr_interface.launch()