Update app.py

app.py

@@ -77,11 +77,17 @@ COLOR_MAPPING_RGB = {to_rgb(k): v for k, v in COLOR_MAPPING_.items()}
 def map_colors_rgb(color: tuple) -> str:
     return COLOR_MAPPING_RGB[color]
 
-def get_segmentation_of_room(image: Image) -> tuple[np.ndarray, Image]:
+def get_segmentation_of_room(image: Image) -> tuple[np.ndarray, Image.Image]:
+
+    # Ensure image is valid before processing
+    if not isinstance(image, Image.Image):
+        raise TypeError("Input 'image' must be a PIL Image object.")
+
+        
     # Semantic Segmentation
     with torch.inference_mode():
-        semantic_inputs = processor(images=image, task_inputs=["semantic"], return_tensors="pt", size={"height": 256, "width": 256})
-        semantic_inputs = {key: value.to("cuda") for key, value in semantic_inputs.items()}
+        semantic_inputs = processor(images=imag, return_tensors="pt", size={"height": 256, "width": 256})
+        semantic_inputs = {key: value.to(DEVICE) for key, value in semantic_inputs.items()}
         semantic_outputs = model(**semantic_inputs)
         # pass through image_processor for postprocessing
         segmentation_maps = processor.post_process_semantic_segmentation(semantic_outputs, target_sizes=[image.size[::-1]])