Update app.py

app.py

@@ -76,19 +76,25 @@ tokenizer = MBartTokenizer.from_pretrained(
     'facebook/mbart-large-50'
 )
 processortext2 = CustomOCRProcessor(image_processor,tokenizer)
-import os
 
+import os
 huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")
-model2 = VisionEncoderDecoderModel.from_pretrained("musadac/vilanocr-single-urdu", use_auth_token=huggingface_token)
+model = {}
+model['single-urdu'] = VisionEncoderDecoderModel.from_pretrained("musadac/vilanocr-single-urdu", use_auth_token=huggingface_token)
+model['multi-urdu'] = VisionEncoderDecoderModel.from_pretrained("musadac/ViLanOCR", use_auth_token=huggingface_token)
+model['medical'] = VisionEncoderDecoderModel.from_pretrained("musadac/vilanocr-multi-medical", use_auth_token=huggingface_token)
+model['chinese'] = VisionEncoderDecoderModel.from_pretrained("musadac/vilanocr-single-chinese", use_auth_token=huggingface_token)
+
 st.title("Image OCR with musadac/vilanocr")
+model_name = st.selectbox("Choose an OCR model", ["single-urdu", "multi-urdu", "medical","chinese" ])
 uploaded_file = st.file_uploader("Choose an image", type=["jpg", "jpeg", "png"])
-
 if uploaded_file is not None:
+    
     img = Image.open(uploaded_file).convert("RGB")
     pixel_values = processortext2(img.convert("RGB"), return_tensors="pt").pixel_values
     
     with torch.no_grad():
-        generated_ids = model2.generate(pixel_values)
+        generated_ids = model[model_name].generate(pixel_values)
     
     result = processortext2.batch_decode(generated_ids, skip_special_tokens=True)[0]
     st.write("OCR Result:")