## 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: | def process_image_to_image() -> None: | ||||||
|     if not modules.globals.nsfw: |     if modules.globals.nsfw: | ||||||
|         from modules.predicter import predict_image |         from modules.predicter import predict_image | ||||||
|         if predict_image(modules.globals.target_path): |         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): |     for frame_processor in get_frame_processors_modules(modules.globals.frame_processors): | ||||||
|         update_status('Processing...', frame_processor.NAME) |         update_status('Processing...', frame_processor.NAME) | ||||||
|         frame_processor.process_image(modules.globals.source_path, modules.globals.output_path, modules.globals.output_path) |         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: | def process_image_to_video() -> None: | ||||||
|     if not modules.globals.nsfw: |     if modules.globals.nsfw: | ||||||
|         from modules.predicter import predict_video |         from modules.predicter import predict_video | ||||||
|         if predict_video(modules.globals.target_path): |         if predict_video(modules.globals.target_path): | ||||||
|             destroy() |             destroy(to_quit=False) | ||||||
|  |             update_status('Processing to video ignored!') | ||||||
|  |             return | ||||||
| 
 | 
 | ||||||
|     update_status('Creating temporary resources...') |     update_status('Creating temporary resources...') | ||||||
|     create_temp(modules.globals.target_path) |     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) |         move_temp(modules.globals.target_path, modules.globals.output_path) | ||||||
| 
 | 
 | ||||||
| 
 | 
 | ||||||
| def destroy() -> None: | def destroy(to_quit=True) -> None: | ||||||
|     if modules.globals.target_path: |     if modules.globals.target_path: | ||||||
|         clean_temp(modules.globals.target_path) |         clean_temp(modules.globals.target_path) | ||||||
|     quit() |     if to_quit: quit() | ||||||
| 
 | 
 | ||||||
| 
 | 
 | ||||||
| def run() -> None: | def run() -> None: | ||||||
|  |  | ||||||
|  | @ -6,9 +6,11 @@ from modules.typing import Frame | ||||||
| MAX_PROBABILITY = 0.85 | MAX_PROBABILITY = 0.85 | ||||||
| 
 | 
 | ||||||
| # Preload the model once for efficiency | # Preload the model once for efficiency | ||||||
| model = opennsfw2.make_open_nsfw_model() | model = None | ||||||
| 
 | 
 | ||||||
| def predict_frame(target_frame: Frame) -> bool: | 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 = Image.fromarray(target_frame) | ||||||
|     image = opennsfw2.preprocess_image(image, opennsfw2.Preprocessing.YAHOO) |     image = opennsfw2.preprocess_image(image, opennsfw2.Preprocessing.YAHOO) | ||||||
|     views = np.expand_dims(image, axis=0) |     views = np.expand_dims(image, axis=0) | ||||||
|  | @ -16,9 +18,9 @@ def predict_frame(target_frame: Frame) -> bool: | ||||||
|     return probability > MAX_PROBABILITY |     return probability > MAX_PROBABILITY | ||||||
| 
 | 
 | ||||||
| def predict_image(target_path: str) -> bool: | 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 |     return probability > MAX_PROBABILITY | ||||||
| 
 | 
 | ||||||
| def predict_video(target_path: str) -> bool: | 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) |     return any(probability > MAX_PROBABILITY for probability in probabilities) | ||||||
|  |  | ||||||
		Loading…
	
		Reference in New Issue