Spaces:
Sleeping
Sleeping
import base64 | |
from datetime import datetime | |
import io | |
import pygame, os | |
import logging | |
import time | |
def capture(pygame_screen, output_dir): | |
if output_dir == "": | |
pygame.image.save(pygame_screen, "tmp.png") | |
filepath = "tmp.png" | |
else: | |
filepath = os.path.join(output_dir, "screen_" + datetime.now().strftime("%Y-%m-%d-%H-%M-%S-%f") + ".png") | |
logging.info("filepath: " + filepath) | |
print("filepath: " + filepath) | |
# 获取当前屏幕的大小 | |
screen_width, screen_height = pygame_screen.get_size() | |
image_size_height = 360 | |
image_size = (int(screen_width * image_size_height / screen_height), image_size_height) | |
print("image_size: ", image_size) | |
# 创建一个新表面,并将当前屏幕内容缩放到该表面 | |
scaled_surface = pygame.transform.scale(pygame_screen, image_size) | |
if os.path.exists(output_dir): | |
# 使用pygame.image.save将缩放后的内容保存到文件 | |
pygame.image.save(scaled_surface, filepath) | |
# pygame.image.save(pygame_screen, filepath) | |
if os.path.exists(filepath): | |
with open(filepath, "rb") as image_file: | |
encoded_string = base64.b64encode(image_file.read()).decode('utf-8') | |
else: | |
byte_io = io.BytesIO() | |
pygame.image.save(pygame_screen, byte_io) | |
byte_io.seek(0) # 将文件指针移到开始 | |
encoded_string = base64.b64encode(byte_io.read()).decode('utf-8') | |
time.sleep(0.02) | |
return filepath, encoded_string | |
def seed_everything(seed): | |
try: | |
import random | |
random.seed(seed) | |
print("random seed set to ", seed) | |
except: | |
print("random seed failed") | |
try: | |
import numpy as np | |
np.random.seed(seed) | |
print("numpy seed set to ", seed) | |
except: | |
print("numpy seed failed") | |
try: | |
import torch | |
torch.manual_seed(seed) | |
if torch.cuda.is_available(): | |
torch.cuda.manual_seed(seed) | |
torch.cuda.manual_seed_all(seed) | |
torch.backends.cudnn.deterministic = True | |
print("torch seed set to ", seed) | |
except: | |
print("torch seed failed") | |