Update app.py

app.py

@@ -42,19 +42,44 @@ print("Loading the transformer model...")
 transformer_vectorizer = TransformerVectorizer()
 vectorizer = TfidfVectorizer()
 def process_input(input_type, user_input, uploaded_file):
-    print('ooooocr')
     if input_type == "File Upload" and uploaded_file is not None:
-
-        with open(uploaded_file.name, "rb") as f:
-            image = f.read()
-        results = reader.readtext(image)
-
-        extracted_text = ' '.join([text[1] for text in results])
-        print("text:")
-        print(extracted_text)
+        file_ext = os.path.splitext(uploaded_file.name)[1].lower()
+        extracted_text = ""
+        
+        if file_ext in ['.jpg', '.jpeg', '.png']:
+            # 处理图片文件
+            results = reader.readtext(uploaded_file.name)
+            extracted_text = ' '.join([text[1] for text in results])
+            print("从图片提取的文本:")
+            print(extracted_text)
+        
+        elif file_ext == '.txt':
+            # 处理TXT文件
+            with open(uploaded_file.name, 'r', encoding='utf-8') as f:
+                extracted_text = f.read()
+            print("从TXT文件提取的文本:")
+            print(extracted_text)
+        
+        elif file_ext == '.pdf':
+            # 处理PDF文件
+            with open(uploaded_file.name, 'rb') as f:
+                reader_pdf = PyPDF2.PdfReader(f)
+                for page_num in range(len(reader_pdf.pages)):
+                    page = reader_pdf.pages[page_num]
+                    extracted_text += page.extract_text() + "\n"
+            print("从PDF文件提取的文本:")
+            print(extracted_text)
+        
+        else:
+            return "不支持的文件类型。请上传 .jpg, .jpeg, .png, .txt 或 .pdf 文件。"
+        
         return extracted_text
+
     elif input_type == "Text Input":
         return user_input
+    else:
+        return "无效的输入类型或未上传文件。
+        
 def toggle_visibility(input_type):
     user_input_visible = input_type == "Text Input"
     file_upload_visible = input_type == "File Upload"