Update app.py

app.py

@@ -11,33 +11,45 @@ def load_model():
     model = AutoModelForImageTextToText.from_pretrained("HuggingFaceTB/SmolVLM-Instruct")
     return processor, model
 
-# Extract text from image
+# Function to preprocess image and handle model execution
 def extract_text(image, processor, model):
     # Initialize progress bar
     progress_bar = st.progress(0)
-    
-    # Simulate steps for progress
-    progress_bar.progress(20)  # Step 1: Starting processing
     time.sleep(0.5)
 
-    # Preprocess image (extract pixel values only)
-    inputs = processor(images=image, return_tensors="pt").to("cpu")
-    pixel_values = inputs.get("pixel_values")
-    progress_bar.progress(50)  # Step 2: Preprocessed image
-    time.sleep(0.5)
+    # Resize the image to fixed dimensions
+    try:
+        required_size = (224, 224)  # Explicit resizing for model input
+        image_resized = image.resize(required_size)
+        progress_bar.progress(20)  # Step 1: Image resized
+        time.sleep(0.5)
+
+        # Preprocess image (extract pixel values)
+        inputs = processor(images=image_resized, return_tensors="pt", do_resize=False).to("cpu")
+        pixel_values = inputs.get("pixel_values")
 
-    # Perform generation
-    with torch.no_grad():
-        outputs = model.generate(pixel_values=pixel_values)
-        progress_bar.progress(80)  # Step 3: Model processing
+        # Check if pixel values are valid
+        if pixel_values is None or pixel_values.shape[0] == 0:
+            raise ValueError("Preprocessing failed: Empty tensor generated for image.")
+
+        progress_bar.progress(50)  # Step 2: Image preprocessed
         time.sleep(0.5)
 
-    # Decode outputs
-    result = processor.batch_decode(outputs, skip_special_tokens=True)[0]
-    progress_bar.progress(100)  # Step 4: Completed
-    time.sleep(0.5)
+        # Perform inference
+        with torch.no_grad():
+            outputs = model.generate(pixel_values=pixel_values)
+            progress_bar.progress(80)  # Step 3: Model processing
+            time.sleep(0.5)
 
-    return result
+        # Decode outputs to text
+        result = processor.batch_decode(outputs, skip_special_tokens=True)[0]
+        progress_bar.progress(100)  # Step 4: Completed
+        time.sleep(0.5)
+
+        return result
+
+    except Exception as e:
+        raise RuntimeError(f"Error during text extraction: {str(e)}")
 
 # Streamlit UI
 def main():
@@ -51,18 +63,19 @@ def main():
     uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])
 
     if uploaded_file is not None:
-        # Open image
-        image = Image.open(uploaded_file).convert("RGB")
-        st.image(image, caption="Uploaded Image", use_container_width=True)
+        try:
+            # Open and display image
+            image = Image.open(uploaded_file).convert("RGB")
+            st.image(image, caption="Uploaded Image", use_container_width=True)
 
-        # Extract text with progress bar
-        with st.spinner("Extracting text... Please wait!"):
-            try:
+            # Extract text with progress bar
+            with st.spinner("Extracting text... Please wait!"):
                 extracted_text = extract_text(image, processor, model)
                 st.subheader("📝 Extracted Text:")
                 st.write(extracted_text)
-            except Exception as e:
-                st.error(f"An error occurred: {str(e)}")
+
+        except Exception as e:
+            st.error(f"An error occurred: {str(e)}")
 
 if __name__ == "__main__":
     main()