Update app.py

app.py

@@ -1,43 +1,25 @@
-import gc
 import torch
 from transformers import TrOCRProcessor, VisionEncoderDecoderModel
-from craft_text_detector import Craft
 from PIL import Image
-import cv2
-import time
 import gradio as gr
 
-# Force CPU usage, disable CUDA
+# Force CPU usage
 torch.set_default_device('cpu')
-craft = Craft(output_dir=None, crop_type="box", cuda=False)
 
-# Load smaller model suitable for CPU
-processor = TrOCRProcessor.from_pretrained('microsoft/trocr-small-handwritten')
-model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-small-handwritten')
+# Load model and processor
+processor = TrOCRProcessor.from_pretrained('microsoft/trocr-base-handwritten')
+model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-base-handwritten')
 
 def recognize_handwritten(image):
-    start_time = time.time()
+    # Convert uploaded image to RGB
+    image = image.convert("RGB")
+    pixel_values = processor(images=image, return_tensors="pt").pixel_values
     
-    # Convert Gradio image to OpenCV format
-    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-    result = craft.detect_text(image=image)
-    boxes = result["boxes"]
-    pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
-    texts = []
+    # Generate text
+    generated_ids = model.generate(pixel_values)
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
     
-    for box in boxes:
-        crop = pil_image.crop([box[0][0], box[0][1], box[2][0], box[2][1]])
-        pixel_values = processor(crop, return_tensors="pt").pixel_values
-        with torch.no_grad():
-            generated_ids = model.generate(pixel_values)
-        text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-        texts.append(text)
-    
-    text_data = " ".join(texts)
-    end_time = time.time()
-    time_difference = end_time - start_time
-    
-    return f"Recognized text: {text_data}\nTime: {time_difference} seconds"
+    return f"Recognized text: {generated_text}"
 
 # Create Gradio interface
 interface = gr.Interface(
@@ -49,8 +31,4 @@ interface = gr.Interface(
 )
 
 # Launch the app
-interface.launch()
-
-# Cleanup
-craft.unload_craftnet_model()
-gc.collect()
+interface.launch()