Choose correct driver to render to the framebuffer

https://learn.adafruit.com/pi-video-output-using-pygame/pointing-pygame-to-the-framebuffer

photoboite.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 
-import sys, pygame
+import os, sys, pygame
 
 class Button:
     def __init__(self, size, topleft):
@@ -46,17 +46,12 @@ class Photoboite:
     i = 0
 
     def __init__(self):
-        pygame.init()
+        self.screen = self.screen()
+        pygame.font.init()
 
-        self.size = 820, 540
-        infos = pygame.display.Info()
-        # self.size = infos.current_w, infos.current_h
-        self.size = 1280, 768
         self.background = 255, 255, 255
         self.buttons = []
 
-        self.screen = pygame.display.set_mode(self.size)
-
         button_size = int(self.screen.get_width() / 4)
         topleft = self.screen.get_width() - button_size - 20, self.screen.get_height() / 4 - button_size / 2
 
@@ -68,6 +63,30 @@ class Photoboite:
         self.font = pygame.font.Font(None, 80)
 
 
+    def screen(self):
+        if os.getenv('DISPLAY'):
+            pygame.display.init()
+            return pygame.display.set_mode((1280, 768))
+        else:
+            drivers = ['directfb', 'fbcon', 'svgalib']
+            found = False
+            for driver in drivers:
+                if not os.getenv('SDL_VIDEODRIVER'):
+                    os.putenv('SDL_VIDEODRIVER', driver)
+                try:
+                    pygame.display.init()
+                except pygame.error as e:
+                    print('Driver: {0} failed.'.format(driver))
+                    continue
+                found = True
+                break
+
+            if not found:
+               raise Exception('No suitable video driver found!')
+
+            return pygame.display.set_mode()
+
+
     def run(self):
         pressed = False
         while True: