## Fix & Improve the NSFW function
- Fixed incorrect state usage. - Removed the redundant argument that caused exceptions. - Prevented the app from closing when an image is flagged as NSFW.pull/237/head
							parent
							
								
									78c808ef36
								
							
						
					
					
						commit
						408b0f4cf7
					
				|  | @ -222,12 +222,18 @@ def start() -> None: | |||
| 
 | ||||
| 
 | ||||
| def process_image_to_image() -> None: | ||||
|     if not modules.globals.nsfw: | ||||
|     if modules.globals.nsfw: | ||||
|         from modules.predicter import predict_image | ||||
|         if predict_image(modules.globals.target_path): | ||||
|             destroy() | ||||
|             destroy(to_quit=False) | ||||
|             update_status('Processing to image ignored!') | ||||
|             return | ||||
| 
 | ||||
|     try: | ||||
|         shutil.copy2(modules.globals.target_path, modules.globals.output_path) | ||||
|     except Exception as e: | ||||
|         print("Error copying file:", str(e)) | ||||
| 
 | ||||
|     shutil.copy2(modules.globals.target_path, modules.globals.output_path) | ||||
|     for frame_processor in get_frame_processors_modules(modules.globals.frame_processors): | ||||
|         update_status('Processing...', frame_processor.NAME) | ||||
|         frame_processor.process_image(modules.globals.source_path, modules.globals.output_path, modules.globals.output_path) | ||||
|  | @ -240,10 +246,12 @@ def process_image_to_image() -> None: | |||
| 
 | ||||
| 
 | ||||
| def process_image_to_video() -> None: | ||||
|     if not modules.globals.nsfw: | ||||
|     if modules.globals.nsfw: | ||||
|         from modules.predicter import predict_video | ||||
|         if predict_video(modules.globals.target_path): | ||||
|             destroy() | ||||
|             destroy(to_quit=False) | ||||
|             update_status('Processing to video ignored!') | ||||
|             return | ||||
| 
 | ||||
|     update_status('Creating temporary resources...') | ||||
|     create_temp(modules.globals.target_path) | ||||
|  | @ -287,10 +295,10 @@ def handle_video_audio() -> None: | |||
|         move_temp(modules.globals.target_path, modules.globals.output_path) | ||||
| 
 | ||||
| 
 | ||||
| def destroy() -> None: | ||||
| def destroy(to_quit=True) -> None: | ||||
|     if modules.globals.target_path: | ||||
|         clean_temp(modules.globals.target_path) | ||||
|     quit() | ||||
|     if to_quit: quit() | ||||
| 
 | ||||
| 
 | ||||
| def run() -> None: | ||||
|  |  | |||
|  | @ -6,9 +6,11 @@ from modules.typing import Frame | |||
| MAX_PROBABILITY = 0.85 | ||||
| 
 | ||||
| # Preload the model once for efficiency | ||||
| model = opennsfw2.make_open_nsfw_model() | ||||
| model = None | ||||
| 
 | ||||
| def predict_frame(target_frame: Frame) -> bool: | ||||
|     global model | ||||
|     if model is None: model = opennsfw2.make_open_nsfw_model() | ||||
|     image = Image.fromarray(target_frame) | ||||
|     image = opennsfw2.preprocess_image(image, opennsfw2.Preprocessing.YAHOO) | ||||
|     views = np.expand_dims(image, axis=0) | ||||
|  | @ -16,9 +18,9 @@ def predict_frame(target_frame: Frame) -> bool: | |||
|     return probability > MAX_PROBABILITY | ||||
| 
 | ||||
| def predict_image(target_path: str) -> bool: | ||||
|     probability = opennsfw2.predict_image(target_path, model=model) | ||||
|     probability = opennsfw2.predict_image(target_path) | ||||
|     return probability > MAX_PROBABILITY | ||||
| 
 | ||||
| def predict_video(target_path: str) -> bool: | ||||
|     _, probabilities = opennsfw2.predict_video_frames(video_path=target_path, frame_interval=100, model=model) | ||||
|     _, probabilities = opennsfw2.predict_video_frames(video_path=target_path, frame_interval=100) | ||||
|     return any(probability > MAX_PROBABILITY for probability in probabilities) | ||||
|  |  | |||
		Loading…
	
		Reference in New Issue